From a0837b391c219858a732043368eb23f9bbbdf7fc Mon Sep 17 00:00:00 2001 From: kouzhenqi <826563886@qq.com> Date: Sat, 18 Jul 2026 10:28:01 +0800 Subject: [PATCH 001/122] =?UTF-8?q?fix(openai):=20=E4=BF=AE=E5=A4=8D=20GPT?= =?UTF-8?q?=20Image=20=E7=94=9F=E6=88=90=E7=A8=B3=E5=AE=9A=E6=80=A7?= =?UTF-8?q?=E5=B9=B6=E5=8F=91=E5=B8=83=201.2.7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 为 Images 非流式请求增加独立总超时，并将 OpenAI 响应头默认超时调整为 600 秒 - 重构 OAuth 图片 SSE 增量读取、心跳、客户端断开排空与首输出前故障转移 - 完善图片参数、上传大小、映射模型与 Responses 图片模型协议校验 - 区分请求错误和账号能力错误，优化账号及备用分组调度与健康状态统计 - 延长图片账号连通性测试并避免瞬态错误永久标记账号异常 - 更新 1.2.7 版本号、示例配置与文档站更新日志 - 验证 go test ./internal/...、go vet、服务端编译、前端构建和文档站构建 --- backend/cmd/server/VERSION | 2 +- backend/internal/config/config.go | 12 +- backend/internal/config/config_test.go | 20 + backend/internal/handler/gateway_handler.go | 10 +- .../handler/openai_gateway_handler_test.go | 145 +++++ backend/internal/handler/openai_images.go | 64 +- .../repository/account_share_mode_repo.go | 6 + backend/internal/service/account.go | 4 +- .../internal/service/account_share_mode.go | 114 ++-- .../internal/service/account_test_service.go | 6 +- .../account_test_service_openai_image_test.go | 2 + backend/internal/service/gateway_service.go | 9 +- .../service/openai_account_scheduler.go | 16 +- .../service/openai_gateway_service.go | 63 ++ backend/internal/service/openai_images.go | 545 ++++++++++++++++-- .../service/openai_images_responses.go | 240 +++++--- .../internal/service/openai_images_test.go | 355 +++++++++++- .../service/openai_oauth_passthrough_test.go | 229 ++++++++ deploy/config.example.yaml | 12 +- .../content/docs/operations/changelog.mdx | 10 + 20 files changed, 1666 insertions(+), 198 deletions(-) diff --git a/backend/cmd/server/VERSION b/backend/cmd/server/VERSION index 3c43790f5..c04c650a7 100644 --- a/backend/cmd/server/VERSION +++ b/backend/cmd/server/VERSION @@ -1 +1 @@ -1.2.6 +1.2.7 diff --git a/backend/internal/config/config.go b/backend/internal/config/config.go index 3c503e36b..17f7e5d8e 100644 --- a/backend/internal/config/config.go +++ b/backend/internal/config/config.go @@ -772,8 +772,11 @@ type GatewayConfig struct { // 注意：这不影响流式数据传输，只控制等待响应头的时间 ResponseHeaderTimeout int `mapstructure:"response_header_timeout"` // OpenAIResponseHeaderTimeout: OpenAI/Codex 上游等待响应头的超时时间（秒），0表示无超时。 - // OpenAI/Codex 请求可能在上游排队较久；默认不使用通用响应头超时截断。 + // OpenAI/Codex 请求可能在上游排队较久，因此使用独立且更宽松的超时。 OpenAIResponseHeaderTimeout int `mapstructure:"openai_response_header_timeout"` + // ImageNonstreamTotalTimeoutSeconds: Images 非流式请求的总超时时间（秒），0表示禁用。 + // 图片生成可能长时间没有响应体数据，不能复用普通流数据间隔超时。 + ImageNonstreamTotalTimeoutSeconds int `mapstructure:"image_nonstream_total_timeout_seconds"` // OpenAIFirstOutputTimeoutSeconds: OpenAI 原生 HTTP Responses 首个语义输出超时（秒），0 表示禁用。 OpenAIFirstOutputTimeoutSeconds int `mapstructure:"openai_first_output_timeout_seconds"` // OpenAIHighEffortFirstOutputTimeoutSeconds: high/xhigh/max 推理的首个语义输出超时（秒）。 @@ -1987,7 +1990,8 @@ func setDefaults() { // Gateway viper.SetDefault("gateway.response_header_timeout", 600) // 600秒(10分钟)等待上游响应头，LLM高负载时可能排队较久 - viper.SetDefault("gateway.openai_response_header_timeout", 0) + viper.SetDefault("gateway.openai_response_header_timeout", 600) + viper.SetDefault("gateway.image_nonstream_total_timeout_seconds", 1800) viper.SetDefault("gateway.openai_first_output_timeout_seconds", 0) viper.SetDefault("gateway.openai_high_effort_first_output_timeout_seconds", 0) viper.SetDefault("gateway.log_upstream_error_body", true) @@ -2713,6 +2717,10 @@ func (c *Config) Validate() error { if c.Gateway.OpenAIResponseHeaderTimeout < 0 { return fmt.Errorf("gateway.openai_response_header_timeout must be non-negative") } + if c.Gateway.ImageNonstreamTotalTimeoutSeconds < 0 || c.Gateway.ImageNonstreamTotalTimeoutSeconds > 3600 || + (c.Gateway.ImageNonstreamTotalTimeoutSeconds > 0 && c.Gateway.ImageNonstreamTotalTimeoutSeconds < 60) { + return fmt.Errorf("gateway.image_nonstream_total_timeout_seconds must be 0 or between 60-3600 seconds") + } if c.Gateway.OpenAIFirstOutputTimeoutSeconds < 0 || c.Gateway.OpenAIFirstOutputTimeoutSeconds > 600 || (c.Gateway.OpenAIFirstOutputTimeoutSeconds > 0 && c.Gateway.OpenAIFirstOutputTimeoutSeconds < 30) { return fmt.Errorf("gateway.openai_first_output_timeout_seconds must be 0 or between 30-600 seconds") diff --git a/backend/internal/config/config_test.go b/backend/internal/config/config_test.go index c25b6b55a..69ed394eb 100644 --- a/backend/internal/config/config_test.go +++ b/backend/internal/config/config_test.go @@ -144,6 +144,16 @@ func TestLoadDefaultSchedulingConfig(t *testing.T) { } } +func TestLoadDefaultOpenAIImageTimeoutConfig(t *testing.T) { + resetViperWithJWTSecret(t) + + cfg, err := Load() + require.NoError(t, err) + require.Equal(t, 600, cfg.Gateway.OpenAIResponseHeaderTimeout) + require.Equal(t, 1800, cfg.Gateway.ImageNonstreamTotalTimeoutSeconds) + require.Equal(t, DefaultUpstreamResponseReadMaxBytes, cfg.Gateway.UpstreamResponseReadMaxBytes) +} + func TestLoadDefaultOpenAIWSConfig(t *testing.T) { resetViperWithJWTSecret(t) @@ -1543,6 +1553,16 @@ func TestValidateConfigErrors(t *testing.T) { mutate: func(c *Config) { c.Gateway.StreamKeepaliveInterval = 4 }, wantErr: "gateway.stream_keepalive_interval", }, + { + name: "gateway image nonstream total timeout negative", + mutate: func(c *Config) { c.Gateway.ImageNonstreamTotalTimeoutSeconds = -1 }, + wantErr: "gateway.image_nonstream_total_timeout_seconds must be 0 or between 60-3600 seconds", + }, + { + name: "gateway image nonstream total timeout range", + mutate: func(c *Config) { c.Gateway.ImageNonstreamTotalTimeoutSeconds = 59 }, + wantErr: "gateway.image_nonstream_total_timeout_seconds must be 0 or between 60-3600 seconds", + }, { name: "gateway image nonstream keepalive negative", mutate: func(c *Config) { c.Gateway.ImageNonstreamKeepaliveInterval = -1 }, diff --git a/backend/internal/handler/gateway_handler.go b/backend/internal/handler/gateway_handler.go index c0ec8b072..ebbca799b 100644 --- a/backend/internal/handler/gateway_handler.go +++ b/backend/internal/handler/gateway_handler.go @@ -1325,7 +1325,11 @@ func (c *apiKeyGroupRouteCursor) recordSuccess(apiKeyID int64) { } func canSwitchAPIKeyGroupRouteAfterForward(c *gin.Context, cursor *apiKeyGroupRouteCursor, failoverErr *service.UpstreamFailoverError, streamStarted bool, writerSizeBeforeForward int) bool { - if cursor == nil || !cursor.hasNext() || !shouldSwitchAPIKeyGroupRoute(failoverErr) || streamStarted { + if cursor == nil || !cursor.hasNext() || !shouldSwitchAPIKeyGroupRoute(failoverErr) { + return false + } + safeAfterWrite := failoverErr != nil && failoverErr.SafeToFailoverAfterWrite + if streamStarted && !safeAfterWrite { return false } if c != nil && c.Writer != nil { @@ -1333,7 +1337,7 @@ func canSwitchAPIKeyGroupRouteAfterForward(c *gin.Context, cursor *apiKeyGroupRo if service.OpenAIImagesJSONKeepalivePresent(c) { writtenSize = service.OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c) } - if writtenSize != writerSizeBeforeForward { + if writtenSize != writerSizeBeforeForward && !safeAfterWrite { return false } } @@ -1414,6 +1418,8 @@ func shouldSwitchAPIKeyGroupRoute(failoverErr *service.UpstreamFailoverError) bo return false } switch failoverErr.StatusCode { + case http.StatusBadRequest: + return failoverErr.Scope == service.GatewayFailureScopeAccount case http.StatusTooManyRequests, http.StatusBadGateway, http.StatusServiceUnavailable, http.StatusGatewayTimeout, 529: return true default: diff --git a/backend/internal/handler/openai_gateway_handler_test.go b/backend/internal/handler/openai_gateway_handler_test.go index 646e8e905..e35c21a6d 100644 --- a/backend/internal/handler/openai_gateway_handler_test.go +++ b/backend/internal/handler/openai_gateway_handler_test.go @@ -128,6 +128,68 @@ func TestOpenAIEnsureForwardErrorResponseDoesNotAppendSecondImageJSON(t *testing require.Equal(t, originalBody, recorder.Body.String()) } +func TestOpenAIImagesForwardMayFailoverOnlyBeforeSemanticWriteOrWhenExplicitlySafe(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + writtenBefore := service.OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c) + + require.True(t, openAIImagesForwardMayFailover(c, writtenBefore, nil)) + + _, err := c.Writer.Write([]byte("semantic-output")) + require.NoError(t, err) + require.False(t, openAIImagesForwardMayFailover(c, writtenBefore, nil)) + require.False(t, openAIImagesForwardMayFailover(c, writtenBefore, &service.UpstreamFailoverError{})) + require.True(t, openAIImagesForwardMayFailover(c, writtenBefore, &service.UpstreamFailoverError{SafeToFailoverAfterWrite: true})) +} + +func TestOpenAIImagesForwardMayFailoverAfterJSONKeepalivePadding(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil) + stop := service.StartOpenAIImagesJSONKeepalive(c, time.Millisecond) + defer stop() + writtenBefore := service.OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c) + + require.Eventually(t, c.Writer.Written, time.Second, time.Millisecond) + require.Equal(t, writtenBefore, service.OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c)) + require.True(t, openAIImagesForwardMayFailover(c, writtenBefore, nil)) +} + +func TestOpenAIImagesRequestFailureIsNotReportedAsAccountFailure(t *testing.T) { + requestErr := &service.UpstreamFailoverError{ + Scope: service.GatewayFailureScopeRequest, + NextAccountAction: service.NextAccountStop, + ClientStatusCode: http.StatusBadRequest, + ClientMessage: "n is not supported for this account route", + } + require.False(t, shouldReportOpenAIImagesScheduleFailure(requestErr)) + + accountErr := &service.UpstreamFailoverError{ + Scope: service.GatewayFailureScopeAccount, + NextAccountAction: service.NextAccountRetry, + } + require.True(t, shouldReportOpenAIImagesScheduleFailure(accountErr)) + require.False(t, shouldReportOpenAIImagesScheduleFailure(nil)) +} + +func TestOpenAIImagesRequestFailureReturnsAccurateClientStatus(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil) + h := &OpenAIGatewayHandler{} + + h.handleImagesFailoverExhausted(c, &service.UpstreamFailoverError{ + Scope: service.GatewayFailureScopeRequest, + NextAccountAction: service.NextAccountStop, + ClientStatusCode: http.StatusUnprocessableEntity, + ClientMessage: "unsupported image option", + }, false) + + require.Equal(t, http.StatusUnprocessableEntity, recorder.Code) + require.Equal(t, "invalid_request_error", gjson.Get(recorder.Body.String(), "error.type").String()) + require.Equal(t, "unsupported image option", gjson.Get(recorder.Body.String(), "error.message").String()) +} + func TestAppendOpenAIProxyLogFields(t *testing.T) { base := []zap.Field{zap.Int64("account_id", 7)} @@ -444,6 +506,20 @@ func TestResolveOpenAIForwardDefaultMappedModel(t *testing.T) { }) } +func TestResolveOpenAIAccountSelectionModel(t *testing.T) { + require.Equal(t, "gpt-image-2", resolveOpenAIAccountSelectionModel(" gpt-image-1 ", service.ChannelMappingResult{ + Mapped: true, + MappedModel: " gpt-image-2 ", + })) + require.Equal(t, "gpt-image-1", resolveOpenAIAccountSelectionModel(" gpt-image-1 ", service.ChannelMappingResult{ + Mapped: true, + MappedModel: " ", + })) + require.Equal(t, "gpt-image-1", resolveOpenAIAccountSelectionModel(" gpt-image-1 ", service.ChannelMappingResult{ + MappedModel: "gpt-image-2", + })) +} + func TestResolveOpenAIMessagesDispatchMappedModel(t *testing.T) { t.Run("exact_claude_model_override_wins", func(t *testing.T) { apiKey := &service.APIKey{ @@ -879,6 +955,75 @@ func TestOpenAIForwardMayFailoverOnlyBeforeSemanticWriteOrWhenExplicitlySafe(t * require.True(t, openAIForwardMayFailover(c, writtenBefore, &service.UpstreamFailoverError{SafeToFailoverAfterWrite: true})) } +func TestCanSwitchAPIKeyGroupRouteAfterForwardAllowsOnlyExplicitlySafeWrites(t *testing.T) { + newCursor := func(hasNext bool) *apiKeyGroupRouteCursor { + candidateCount := 1 + if hasNext { + candidateCount = 2 + } + candidates := make([]apiKeyGroupRouteCandidate, candidateCount) + for i := range candidates { + candidates[i].APIKey = &service.APIKey{ID: int64(i + 1)} + } + return newAPIKeyGroupRouteCursorFromCandidates(candidates, true) + } + newContextWithComment := func() (*gin.Context, int) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + writtenBefore := c.Writer.Size() + _, err := c.Writer.Write([]byte(":\n\n")) + require.NoError(t, err) + return c, writtenBefore + } + + t.Run("safe comment permits route switch after stream started", func(t *testing.T) { + c, writtenBefore := newContextWithComment() + failoverErr := &service.UpstreamFailoverError{ + StatusCode: http.StatusBadGateway, + SafeToFailoverAfterWrite: true, + } + require.True(t, canSwitchAPIKeyGroupRouteAfterForward(c, newCursor(true), failoverErr, true, writtenBefore)) + }) + + t.Run("non-safe write remains blocked", func(t *testing.T) { + c, writtenBefore := newContextWithComment() + failoverErr := &service.UpstreamFailoverError{StatusCode: http.StatusBadGateway} + require.False(t, canSwitchAPIKeyGroupRouteAfterForward(c, newCursor(true), failoverErr, true, writtenBefore)) + }) + + t.Run("safe write still requires another route", func(t *testing.T) { + c, writtenBefore := newContextWithComment() + failoverErr := &service.UpstreamFailoverError{ + StatusCode: http.StatusBadGateway, + SafeToFailoverAfterWrite: true, + } + require.False(t, canSwitchAPIKeyGroupRouteAfterForward(c, newCursor(false), failoverErr, true, writtenBefore)) + }) + + t.Run("safe write still requires a route-switchable status", func(t *testing.T) { + c, writtenBefore := newContextWithComment() + failoverErr := &service.UpstreamFailoverError{ + StatusCode: http.StatusBadRequest, + SafeToFailoverAfterWrite: true, + } + require.False(t, canSwitchAPIKeyGroupRouteAfterForward(c, newCursor(true), failoverErr, true, writtenBefore)) + }) +} + +func TestShouldSwitchAPIKeyGroupRouteAllowsOnlyAccountScopedBadRequest(t *testing.T) { + require.True(t, shouldSwitchAPIKeyGroupRoute(&service.UpstreamFailoverError{ + StatusCode: http.StatusBadRequest, + Scope: service.GatewayFailureScopeAccount, + })) + require.False(t, shouldSwitchAPIKeyGroupRoute(&service.UpstreamFailoverError{ + StatusCode: http.StatusBadRequest, + Scope: service.GatewayFailureScopeRequest, + })) + require.False(t, shouldSwitchAPIKeyGroupRoute(&service.UpstreamFailoverError{ + StatusCode: http.StatusBadRequest, + })) +} + func TestOpenAIFirstOutputFailoverExhaustedAllowsOnlyOneAccountSwitch(t *testing.T) { failoverErr := &service.UpstreamFailoverError{SafeToFailoverAfterWrite: true} switchCount := 0 diff --git a/backend/internal/handler/openai_images.go b/backend/internal/handler/openai_images.go index 762e38248..8436f6e00 100644 --- a/backend/internal/handler/openai_images.go +++ b/backend/internal/handler/openai_images.go @@ -138,6 +138,7 @@ routeLoop: return } channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(c.Request.Context(), currentAPIKey.GroupID, parsed.Model) + selectionModel := resolveOpenAIAccountSelectionModel(parsed.Model, channelMapping) if err := h.billingCacheService.CheckBillingEligibility(c.Request.Context(), currentAPIKey.User, currentAPIKey, currentAPIKey.Group, currentSubscription); err != nil { reqLog.Info("openai.images.billing_eligibility_check_failed", zap.Error(err), @@ -173,7 +174,7 @@ routeLoop: selectionCtx, currentAPIKey.GroupID, sessionHash, - parsed.Model, + selectionModel, failedAccountIDs, parsed.RequiredCapability, ) @@ -192,7 +193,7 @@ routeLoop: if routeCursor.switchToNext(apiKey.ID, "account_select_failed", reqLog, zap.Error(err)) { continue routeLoop } - cls := classifyNoAccountErrorFromGin(c, h.gatewayService, currentAPIKey, parsed.Model, parsed.Model, service.PlatformOpenAI) + cls := classifyNoAccountErrorFromGin(c, h.gatewayService, currentAPIKey, selectionModel, parsed.Model, service.PlatformOpenAI) h.handleStreamingAwareError(c, cls.Status, cls.ErrType, cls.Message, streamStarted) return } @@ -201,14 +202,14 @@ routeLoop: routeCursor.switchToNext(apiKey.ID, "account_selection_exhausted", reqLog, zap.Int("upstream_status", lastFailoverErr.StatusCode)) { continue routeLoop } - h.handleFailoverExhausted(c, lastFailoverErr, streamStarted) + h.handleImagesFailoverExhausted(c, lastFailoverErr, streamStarted) } else { h.handleFailoverExhaustedSimple(c, 502, streamStarted) } return } if selection == nil || selection.Account == nil { - cls := classifyNoAccountErrorFromGin(c, h.gatewayService, currentAPIKey, parsed.Model, parsed.Model, service.PlatformOpenAI) + cls := classifyNoAccountErrorFromGin(c, h.gatewayService, currentAPIKey, selectionModel, parsed.Model, service.PlatformOpenAI) h.handleStreamingAwareError(c, cls.Status, cls.ErrType, cls.Message, streamStarted) return } @@ -267,19 +268,22 @@ routeLoop: if failoverClientGone(c) { return } - if failoverErr.ShouldReportAccountScheduleFailure() { + if shouldReportOpenAIImagesScheduleFailure(failoverErr) { h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil) } - if service.OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c) != writerSizeBeforeForward { + if !openAIImagesForwardMayFailover(c, writerSizeBeforeForward, failoverErr) { reqLog.Warn("openai.images.upstream_failover_skipped_after_flush", zap.Int64("account_id", account.ID), zap.Int("upstream_status", failoverErr.StatusCode), ) - h.handleFailoverExhausted(c, failoverErr, true) + h.handleImagesFailoverExhausted(c, failoverErr, true) return } + if failoverErr.SafeToFailoverAfterWrite && c.Writer.Written() && !service.OpenAIImagesJSONKeepalivePresent(c) { + streamStarted = true + } if !failoverErr.ShouldRetryNextAccount() { - h.handleFailoverExhausted(c, failoverErr, streamStarted) + h.handleImagesFailoverExhausted(c, failoverErr, streamStarted) return } if failoverErr.RetryableOnSameAccount { @@ -308,7 +312,7 @@ routeLoop: routeCursor.switchToNext(apiKey.ID, "upstream_failover_exhausted", reqLog, zap.Int("upstream_status", failoverErr.StatusCode)) { continue routeLoop } - h.handleFailoverExhausted(c, failoverErr, streamStarted) + h.handleImagesFailoverExhausted(c, failoverErr, streamStarted) return } switchCount++ @@ -347,6 +351,8 @@ routeLoop: userAgent := c.GetHeader("User-Agent") clientIP := ip.GetClientIP(c) + inboundEndpoint := GetInboundEndpoint(c) + upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform) requestPayloadHash := service.HashUsageRequestPayload(body) if parsed.Multipart { requestPayloadHash = service.HashUsageRequestPayload([]byte(parsed.StickySessionSeed())) @@ -360,8 +366,8 @@ routeLoop: User: currentAPIKey.User, Account: account, Subscription: currentSubscription, - InboundEndpoint: GetInboundEndpoint(c), - UpstreamEndpoint: GetUpstreamEndpoint(c, account.Platform), + InboundEndpoint: inboundEndpoint, + UpstreamEndpoint: upstreamEndpoint, UserAgent: userAgent, IPAddress: clientIP, RequestPayloadHash: requestPayloadHash, @@ -388,6 +394,42 @@ routeLoop: } } +func shouldReportOpenAIImagesScheduleFailure(failoverErr *service.UpstreamFailoverError) bool { + if failoverErr == nil { + return false + } + if failoverErr.Scope == service.GatewayFailureScopeRequest && failoverErr.NextAccountAction == service.NextAccountStop { + return false + } + return failoverErr.ShouldReportAccountScheduleFailure() +} + +func openAIImagesForwardMayFailover(c *gin.Context, writerSizeBeforeForward int, failoverErr *service.UpstreamFailoverError) bool { + if c == nil || c.Writer == nil { + return false + } + if service.OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c) == writerSizeBeforeForward { + return true + } + return failoverErr != nil && failoverErr.SafeToFailoverAfterWrite +} + +func (h *OpenAIGatewayHandler) handleImagesFailoverExhausted(c *gin.Context, failoverErr *service.UpstreamFailoverError, streamStarted bool) { + if failoverErr != nil && + failoverErr.Scope == service.GatewayFailureScopeRequest && + failoverErr.NextAccountAction == service.NextAccountStop && + failoverErr.ClientStatusCode >= http.StatusBadRequest && + failoverErr.ClientStatusCode < http.StatusInternalServerError { + message := strings.TrimSpace(failoverErr.ClientMessage) + if message == "" { + message = http.StatusText(failoverErr.ClientStatusCode) + } + h.handleStreamingAwareError(c, failoverErr.ClientStatusCode, "invalid_request_error", message, streamStarted) + return + } + h.handleFailoverExhausted(c, failoverErr, streamStarted) +} + func (h *OpenAIGatewayHandler) openAIImagesJSONKeepaliveInterval() time.Duration { if h == nil || h.cfg == nil || h.cfg.Gateway.ImageNonstreamKeepaliveInterval <= 0 { return 0 diff --git a/backend/internal/repository/account_share_mode_repo.go b/backend/internal/repository/account_share_mode_repo.go index 7b7c060a5..cba21773a 100644 --- a/backend/internal/repository/account_share_mode_repo.go +++ b/backend/internal/repository/account_share_mode_repo.go @@ -4605,6 +4605,12 @@ func (r *accountShareModeRepository) SuspendMembershipForDispatchFailure(ctx con if err != nil { return nil, err } + // Slot acquisition and its heartbeat update last_request_at only after the + // membership is actually in use. Keep this check inside the membership row + // lock so a concurrent dispatch failure cannot queue an active stream. + if accountShareMembershipRecentlyActive(membership, failedAt) { + return nil, nil + } membership, err = r.suspendActiveMembershipInTx(ctx, tx, membership, failedAt, cooldownUntil) if err != nil { return nil, err diff --git a/backend/internal/service/account.go b/backend/internal/service/account.go index c02acf3d9..c8a717d12 100644 --- a/backend/internal/service/account.go +++ b/backend/internal/service/account.go @@ -2126,8 +2126,10 @@ func (a *Account) SupportsOpenAIImageCapability(capability OpenAIImagesCapabilit return false } switch capability { - case OpenAIImagesCapabilityBasic, OpenAIImagesCapabilityNative: + case OpenAIImagesCapabilityBasic: return a.Type == AccountTypeOAuth || a.Type == AccountTypeAPIKey + case OpenAIImagesCapabilityNative: + return a.Type == AccountTypeAPIKey default: return true } diff --git a/backend/internal/service/account_share_mode.go b/backend/internal/service/account_share_mode.go index 78acaeb89..970dcd466 100644 --- a/backend/internal/service/account_share_mode.go +++ b/backend/internal/service/account_share_mode.go @@ -63,6 +63,8 @@ const ( AccountShareModeEditSessionTTL = 10 * time.Minute AccountShareModeQueueMaxItems = 5 AccountShareModeDispatchCooldown = 5 * time.Minute + AccountShareModeConnectivityTestTimeout = 90 * time.Second + AccountShareModeImageConnectivityTestTimeout = 10 * time.Minute AccountShareRecommendationDefaultLimit = 5 AccountShareRecommendationMaxLimit = 10 AccountShareRecommendationMaxRequests = 1000000 @@ -908,7 +910,6 @@ type AccountShareModeService struct { reviewStopOnce sync.Once reviewStartOnce sync.Once reviewWG sync.WaitGroup - lastRequestTouchL1 sync.Map } func NewAccountShareModeService( @@ -2305,9 +2306,10 @@ func (s *AccountShareModeService) validateOwnerRelist(ctx context.Context, actor return nil } - testCtx, cancel := context.WithTimeout(ctx, 90*time.Second) + modelID := firstAllowedModel(listing.AllowedModels) + testCtx, cancel := context.WithTimeout(ctx, accountShareConnectivityTestTimeout(modelID)) defer cancel() - result, err := s.accountTestService.RunTestBackground(testCtx, listing.AccountID, firstAllowedModel(listing.AllowedModels)) + result, err := s.accountTestService.RunTestBackground(testCtx, listing.AccountID, modelID) if err != nil { return accountShareRelistTestError(err.Error()) } @@ -2345,6 +2347,47 @@ func accountShareRelistTestError(reason string) error { return infraerrors.Newf(http.StatusBadRequest, "ACCOUNT_SHARE_RELIST_TEST_FAILED", "重新上架前自动测试失败：%s", reason) } +func accountShareConnectivityTestTimeout(modelID string) time.Duration { + if isOpenAIImageModel(strings.TrimSpace(modelID)) { + return AccountShareModeImageConnectivityTestTimeout + } + return AccountShareModeConnectivityTestTimeout +} + +func isTransientAccountShareConnectivityFailure(message string) bool { + message = strings.ToLower(strings.TrimSpace(message)) + if message == "" { + return false + } + for _, marker := range []string{ + "context deadline exceeded", + "context canceled", + "request failed", + "failed to read", + "timeout", + "timed out", + "temporary", + "temporarily", + "try again", + "connection reset", + "unexpected eof", + "rate limit", + "rate_limit", + "too many requests", + "overloaded", + "capacity", + "returned 408", + "returned 425", + "returned 429", + "returned 5", + } { + if strings.Contains(message, marker) { + return true + } + } + return false +} + func (s *AccountShareModeService) enrichListingRuntime(ctx context.Context, listing *AccountShareListing) { if listing == nil { return @@ -2793,10 +2836,14 @@ func (s *AccountShareModeService) ResolveActiveBindingForRequest(ctx context.Con } if accountShareListingAccountUnavailableAt(listing, now) { afterRank = membership.QueueRank - result, err := s.suspendMembershipForDispatchFailure(ctx, membership, now) + result, suspended, err := s.suspendMembershipForDispatchFailure(ctx, membership, now) if err != nil { return nil, nil, err } + if !suspended { + lastErr = ErrNoAvailableAccounts + break + } s.invalidateSeatBillingCaches(result) continue } @@ -2808,9 +2855,6 @@ func (s *AccountShareModeService) ResolveActiveBindingForRequest(ctx context.Con afterRank = membership.QueueRank continue } - if err := s.touchMembershipLastRequest(ctx, membership.ID, now); err != nil { - return nil, nil, err - } if requestCtx, ok := AccountShareModeRequestFromContext(ctx); ok && requestCtx.state != nil { requestCtx.state.set(userID, apiKeyID, groupID, membership, listing, nil) } @@ -2876,37 +2920,48 @@ func (s *AccountShareModeService) resolveActiveOrActivateQueuedBinding(ctx conte return membership, listing, nil } -func (s *AccountShareModeService) deferMembershipForDispatchRetry(ctx context.Context, requestCtx AccountShareModeRequestContext, membership *AccountShareMembership, now time.Time) error { +func (s *AccountShareModeService) deferMembershipForDispatchRetry(ctx context.Context, requestCtx AccountShareModeRequestContext, membership *AccountShareMembership, now time.Time) (bool, error) { if s == nil || membership == nil || membership.ID <= 0 { - return nil + return false, nil } - result, err := s.suspendMembershipForDispatchFailure(ctx, membership, now) + result, suspended, err := s.suspendMembershipForDispatchFailure(ctx, membership, now) if err != nil { - return err + return false, err + } + if !suspended { + return false, nil } s.invalidateSeatBillingCaches(result) if requestCtx.state != nil { requestCtx.state.clear() } - return nil + return true, nil } -func (s *AccountShareModeService) suspendMembershipForDispatchFailure(ctx context.Context, membership *AccountShareMembership, now time.Time) (*AccountShareSeatBillingResult, error) { +func (s *AccountShareModeService) suspendMembershipForDispatchFailure(ctx context.Context, membership *AccountShareMembership, now time.Time) (*AccountShareSeatBillingResult, bool, error) { if s == nil || s.repo == nil || membership == nil || membership.ID <= 0 { - return &AccountShareSeatBillingResult{}, nil + return &AccountShareSeatBillingResult{}, false, nil + } + active, err := s.membershipHasActiveConcurrency(ctx, membership.ID) + if err != nil { + return nil, false, err + } + if active { + log.Printf("account_share_mode: dispatch suspension skipped for active membership: membership_id=%d", membership.ID) + return &AccountShareSeatBillingResult{}, false, nil } suspended, err := s.repo.SuspendMembershipForDispatchFailure(ctx, membership.ID, now, now.Add(AccountShareModeDispatchCooldown)) if err != nil { - return nil, err + return nil, false, err } if suspended == nil { - return &AccountShareSeatBillingResult{}, nil + return &AccountShareSeatBillingResult{}, false, nil } return &AccountShareSeatBillingResult{ DebitUserIDs: []int64{suspended.ConsumerUserID}, CreditUserIDs: []int64{suspended.OwnerUserID}, EndedConsumerUserIDs: []int64{suspended.ConsumerUserID}, - }, nil + }, true, nil } func (s *AccountShareModeService) endIdleMembershipForRequest(ctx context.Context, membership *AccountShareMembership, now time.Time) (bool, error) { @@ -2941,23 +2996,6 @@ func (s *AccountShareModeService) endIdleMembershipForRequest(ctx context.Contex return true, nil } -func (s *AccountShareModeService) touchMembershipLastRequest(ctx context.Context, membershipID int64, at time.Time) error { - if s == nil || s.repo == nil || membershipID <= 0 { - return nil - } - now := at.UTC() - if v, ok := s.lastRequestTouchL1.Load(membershipID); ok { - if nextAllowedAt, ok := v.(time.Time); ok && now.Before(nextAllowedAt) { - return nil - } - } - if err := s.repo.TouchMembershipLastRequest(ctx, membershipID, now); err != nil { - return err - } - s.lastRequestTouchL1.Store(membershipID, now.Add(AccountShareModeLastRequestTouchInterval)) - return nil -} - func (s *AccountShareModeService) membershipHasActiveConcurrency(ctx context.Context, membershipID int64) (bool, error) { if s == nil || s.concurrencyService == nil || membershipID <= 0 { return false, nil @@ -3029,7 +3067,7 @@ func (s *AccountShareModeService) schedulePostCreateConnectivityTest(listing *Ac accountID := listing.AccountID modelID := firstAllowedModel(listing.AllowedModels) go func() { - testCtx, cancel := context.WithTimeout(context.Background(), 90*time.Second) + testCtx, cancel := context.WithTimeout(context.Background(), accountShareConnectivityTestTimeout(modelID)) defer cancel() result, err := s.accountTestService.RunTestBackground(testCtx, accountID, modelID) @@ -3046,6 +3084,11 @@ func (s *AccountShareModeService) schedulePostCreateConnectivityTest(listing *Ac if errorMessage == "" { return } + if testCtx.Err() != nil || isTransientAccountShareConnectivityFailure(errorMessage) { + safeMessage := truncateString(sanitizeUpstreamErrorMessage(errorMessage), 512) + log.Printf("account_share_mode: transient post-create connectivity test failure ignored: account_id=%d err=%s", accountID, safeMessage) + return + } writeCtx, writeCancel := context.WithTimeout(context.Background(), 10*time.Second) defer writeCancel() @@ -3127,7 +3170,6 @@ func (s *AccountShareModeService) forceTouchMembershipLastRequest(membershipID i if err := s.repo.TouchMembershipLastRequest(ctx, membershipID, now); err != nil { return err } - s.lastRequestTouchL1.Store(membershipID, now.Add(AccountShareModeLastRequestTouchInterval)) return nil } diff --git a/backend/internal/service/account_test_service.go b/backend/internal/service/account_test_service.go index c0f6b3d18..ca7404d4a 100644 --- a/backend/internal/service/account_test_service.go +++ b/backend/internal/service/account_test_service.go @@ -1644,10 +1644,8 @@ func (s *AccountTestService) testOpenAIImageAPIKey(c *gin.Context, ctx context.C s.sendEvent(c, TestEvent{Type: "test_start", Model: modelID}) payload := map[string]any{ - "model": modelID, - "prompt": prompt, - "n": 1, - "response_format": "b64_json", + "model": modelID, + "prompt": prompt, } payloadBytes, _ := json.Marshal(payload) diff --git a/backend/internal/service/account_test_service_openai_image_test.go b/backend/internal/service/account_test_service_openai_image_test.go index 0f3906159..942bdaf37 100644 --- a/backend/internal/service/account_test_service_openai_image_test.go +++ b/backend/internal/service/account_test_service_openai_image_test.go @@ -88,6 +88,8 @@ func TestAccountTestService_OpenAIImageAPIKeyUsesConfiguredV1BaseURL(t *testing. require.NotNil(t, upstream.lastReq) require.Equal(t, "https://image-upstream.example/v1/images/generations", upstream.lastReq.URL.String()) require.Equal(t, "Bearer test-api-key", upstream.lastReq.Header.Get("Authorization")) + require.NotContains(t, string(upstream.lastBody), "response_format") + require.NotContains(t, string(upstream.lastBody), `"n"`) require.Contains(t, rec.Body.String(), "data:image/png;base64,aGVsbG8=") require.Contains(t, rec.Body.String(), "\"success\":true") } diff --git a/backend/internal/service/gateway_service.go b/backend/internal/service/gateway_service.go index a1ddf1d36..bc0c1bdde 100644 --- a/backend/internal/service/gateway_service.go +++ b/backend/internal/service/gateway_service.go @@ -2323,9 +2323,16 @@ func (s *GatewayService) selectAccountShareModeBoundAccount(ctx context.Context, } if retryCurrentMembership { now := time.Now().UTC() - if err := s.accountShareModeService.deferMembershipForDispatchRetry(ctx, reqCtx, membership, now); err != nil { + deferred, err := s.accountShareModeService.deferMembershipForDispatchRetry(ctx, reqCtx, membership, now) + if err != nil { return nil, true, err } + if !deferred { + if lastErr != nil { + return nil, true, lastErr + } + return nil, true, ErrNoAvailableAccounts + } membership = nil listing = nil account = nil diff --git a/backend/internal/service/openai_account_scheduler.go b/backend/internal/service/openai_account_scheduler.go index d448dbf46..763194167 100644 --- a/backend/internal/service/openai_account_scheduler.go +++ b/backend/internal/service/openai_account_scheduler.go @@ -1338,6 +1338,11 @@ func (s *OpenAIGatewayService) selectAccountShareModeBoundAccount( if !s.accountShareModeService.IsModeGroup(ctx, *groupID) { return nil, decision, false, nil } + boundImageCapability := requiredImageCapability + if boundImageCapability == OpenAIImagesCapabilityNative { + // 共享模式必须保留绑定账号；OAuth 桥接能力由 ForwardImages 做精确参数校验，不能在这里误挂起会员关系。 + boundImageCapability = OpenAIImagesCapabilityBasic + } reqCtx, ok := AccountShareModeRequestFromContext(ctx) if !ok { return nil, decision, true, ErrAccountShareModeGroupUnbound @@ -1402,7 +1407,7 @@ func (s *OpenAIGatewayService) selectAccountShareModeBoundAccount( lastErr = ErrNoAvailableAccounts retryCurrentMembership = true } - if !retryCurrentMembership && !accountSupportsRequestedOpenAIImageCapability(account, requiredImageCapability) { + if !retryCurrentMembership && !accountSupportsRequestedOpenAIImageCapability(account, boundImageCapability) { lastErr = ErrNoAvailableAccounts retryCurrentMembership = true } @@ -1416,9 +1421,16 @@ func (s *OpenAIGatewayService) selectAccountShareModeBoundAccount( } if retryCurrentMembership { now := time.Now().UTC() - if err := s.accountShareModeService.deferMembershipForDispatchRetry(ctx, reqCtx, membership, now); err != nil { + deferred, err := s.accountShareModeService.deferMembershipForDispatchRetry(ctx, reqCtx, membership, now) + if err != nil { return nil, decision, true, err } + if !deferred { + if lastErr != nil { + return nil, decision, true, lastErr + } + return nil, decision, true, ErrNoAvailableAccounts + } membership = nil listing = nil account = nil diff --git a/backend/internal/service/openai_gateway_service.go b/backend/internal/service/openai_gateway_service.go index 00eb1b63c..8cad23fa4 100644 --- a/backend/internal/service/openai_gateway_service.go +++ b/backend/internal/service/openai_gateway_service.go @@ -2508,9 +2508,50 @@ func (s *OpenAIGatewayService) ForwardWithAnalysis(ctx context.Context, c *gin.C reqModel, reqStream, promptCacheKey := analysis.Model, analysis.Stream, analysis.PromptCacheKey originalModel := reqModel + imageOnlyResponsesModel := strings.HasPrefix(strings.ToLower(strings.TrimSpace(reqModel)), "gpt-image-") + rejectImageOnlyResponsesRequest := func(param string, requestErr error) (*OpenAIForwardResult, error) { + setOpsUpstreamError(c, http.StatusBadRequest, requestErr.Error(), "") + c.JSON(http.StatusBadRequest, gin.H{ + "error": gin.H{ + "type": "invalid_request_error", + "message": requestErr.Error(), + "param": param, + }, + }) + return nil, requestErr + } + if imageOnlyResponsesModel { + reqBody, parseErr := getOpenAIRequestBodyMap(c, body) + if parseErr != nil { + return nil, parseErr + } + background := strings.TrimSpace(firstNonEmptyString(reqBody["background"])) + if validationErr := validateOpenAIImagesOptionsForModel(&OpenAIImagesRequest{ + Model: reqModel, + Background: background, + }, reqModel); validationErr != nil { + return rejectImageOnlyResponsesRequest("background", validationErr) + } + if n, exists := reqBody["n"]; exists { + number, validNumber := n.(float64) + if !validNumber || number != 1 { + requestErr := fmt.Errorf("/v1/responses image_generation supports one image per request; n=%v is not supported; use /v1/images/generations for multiple images", n) + return rejectImageOnlyResponsesRequest("n", requestErr) + } + delete(reqBody, "n") + body, err = marshalOpenAIUpstreamJSON(reqBody) + if err != nil { + return nil, fmt.Errorf("remove unsupported /responses image n parameter: %w", err) + } + } + } if account != nil && account.Platform == PlatformGrok { return s.forwardGrokResponses(ctx, c, account, body, originalModel, reqStream, startTime) } + if account.Type == AccountTypeAPIKey && imageOnlyResponsesModel { + requestErr := fmt.Errorf("/v1/responses does not accept image-only model %q as the top-level model for API Key accounts; use /v1/images/generations, or use a Responses-compatible text model with the image_generation tool", reqModel) + return rejectImageOnlyResponsesRequest("model", requestErr) + } if account.Type == AccountTypeAPIKey && !openai_compat.ShouldUseResponsesAPI(account.Extra) { return s.forwardResponsesViaRawChatCompletions(ctx, c, account, body) } @@ -2520,6 +2561,10 @@ func (s *OpenAIGatewayService) ForwardWithAnalysis(ctx context.Context, c *gin.C if isCodexCLI { codexImageGenerationExplicitToolPolicy = account.CodexImageGenerationExplicitToolPolicy() } + if imageOnlyResponsesModel && isCodexCLI && codexImageGenerationExplicitToolPolicy == codexImageGenerationExplicitToolPolicyStrip { + requestErr := fmt.Errorf("/v1/responses image model %q is disabled by this account's Codex image generation policy; use /v1/images/generations or set codex_image_generation_explicit_tool_policy to allow", reqModel) + return rejectImageOnlyResponsesRequest("model", requestErr) + } wsDecision := s.getOpenAIWSProtocolResolver().Resolve(account) clientTransport := GetOpenAIClientTransport(c) // 浠呭厑璁?WS 鍏ョ珯璇锋眰璧?WS 涓婃父锛岄伩鍏嶅嚭鐜?HTTP -> WS 鍗忚娣风敤銆? @@ -2562,6 +2607,24 @@ func (s *OpenAIGatewayService) ForwardWithAnalysis(ctx context.Context, c *gin.C return nil, err } } + if passthroughEnabled && imageOnlyResponsesModel { + reqBody, parseErr := getOpenAIRequestBodyMap(c, body) + if parseErr != nil { + return nil, parseErr + } + if normalizeOpenAIResponsesImageOnlyModel(reqBody) { + body, err = marshalOpenAIUpstreamJSON(reqBody) + if err != nil { + return nil, fmt.Errorf("normalize passthrough image model request: %w", err) + } + logger.LegacyPrintf( + "service.openai_gateway", + "[OpenAI passthrough] Normalized /responses image-only model request inbound_model=%s upstream_model=%s", + reqModel, + openAIImagesResponsesMainModel, + ) + } + } originalBody := body if passthroughEnabled { if isCodexCLI && codexImageGenerationExplicitToolPolicy == codexImageGenerationExplicitToolPolicyStrip { diff --git a/backend/internal/service/openai_images.go b/backend/internal/service/openai_images.go index dc5042724..b58cadbbe 100644 --- a/backend/internal/service/openai_images.go +++ b/backend/internal/service/openai_images.go @@ -11,6 +11,7 @@ import ( "errors" "fmt" "io" + "math" "mime" "mime/multipart" "net/http" @@ -169,6 +170,9 @@ func (s *OpenAIGatewayService) ParseOpenAIImagesRequest(c *gin.Context, body []b if err := validateOpenAIImagesModel(req.Model); err != nil { return nil, err } + if err := validateOpenAIImagesOptions(req); err != nil { + return nil, err + } req.SizeTier = normalizeOpenAIImageSizeTier(req.Size) req.RequiredCapability = classifyOpenAIImagesCapability(req) return req, nil @@ -192,10 +196,11 @@ func parseOpenAIImagesJSONRequest(body []byte, req *OpenAIImagesRequest) error { if nResult.Type != gjson.Number { return fmt.Errorf("invalid n field type") } - req.N = int(nResult.Int()) - if req.N <= 0 { - return fmt.Errorf("n must be greater than 0") + n := nResult.Float() + if n != math.Trunc(n) { + return fmt.Errorf("n must be an integer") } + req.N = int(n) } if sizeResult := gjson.GetBytes(body, "size"); sizeResult.Exists() { @@ -214,14 +219,22 @@ func parseOpenAIImagesJSONRequest(body []byte, req *OpenAIImagesRequest) error { if outputCompression.Type != gjson.Number { return fmt.Errorf("invalid output_compression field type") } - v := int(outputCompression.Int()) + value := outputCompression.Float() + if value != math.Trunc(value) { + return fmt.Errorf("output_compression must be an integer") + } + v := int(value) req.OutputCompression = &v } if partialImages := gjson.GetBytes(body, "partial_images"); partialImages.Exists() { if partialImages.Type != gjson.Number { return fmt.Errorf("invalid partial_images field type") } - v := int(partialImages.Int()) + value := partialImages.Float() + if value != math.Trunc(value) { + return fmt.Errorf("partial_images must be an integer") + } + v := int(value) req.PartialImages = &v } if req.IsEdits() { @@ -282,11 +295,14 @@ func parseOpenAIImagesMultipartRequest(body []byte, contentType string, req *Ope continue } - data, err := io.ReadAll(io.LimitReader(part, openAIImageMaxUploadPartSize)) + data, err := io.ReadAll(io.LimitReader(part, openAIImageMaxUploadPartSize+1)) _ = part.Close() if err != nil { return fmt.Errorf("read multipart field %s: %w", name, err) } + if len(data) > openAIImageMaxUploadPartSize { + return fmt.Errorf("multipart field %s exceeds the 20MB per-part limit", name) + } fileName := strings.TrimSpace(part.FileName()) if fileName != "" { @@ -395,9 +411,6 @@ func applyOpenAIImagesDefaults(req *OpenAIImagesRequest) { if req == nil { return } - if req.N <= 0 { - req.N = 1 - } req.Model = strings.TrimSpace(req.Model) } @@ -424,6 +437,33 @@ func validateOpenAIImagesModel(model string) error { return fmt.Errorf("images endpoint requires an image model, got %q", model) } +func validateOpenAIImagesOptions(req *OpenAIImagesRequest) error { + if req == nil { + return fmt.Errorf("images request is required") + } + if req.N < 1 || req.N > 10 { + return fmt.Errorf("n must be between 1 and 10") + } + if req.PartialImages != nil && (*req.PartialImages < 0 || *req.PartialImages > 3) { + return fmt.Errorf("partial_images must be between 0 and 3") + } + if req.OutputCompression != nil && (*req.OutputCompression < 0 || *req.OutputCompression > 100) { + return fmt.Errorf("output_compression must be between 0 and 100") + } + return validateOpenAIImagesOptionsForModel(req, req.Model) +} + +func validateOpenAIImagesOptionsForModel(req *OpenAIImagesRequest, model string) error { + if req == nil { + return fmt.Errorf("images request is required") + } + if strings.EqualFold(strings.TrimSpace(model), "gpt-image-2") && + strings.EqualFold(strings.TrimSpace(req.Background), "transparent") { + return fmt.Errorf("background transparent is not supported by gpt-image-2") + } + return nil +} + func normalizeOpenAIImagesEndpointPath(path string) string { trimmed := strings.TrimSpace(path) switch { @@ -525,11 +565,14 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesAPIKey( requestModel = mapped } if err := validateOpenAIImagesModel(requestModel); err != nil { - return nil, err + return nil, newOpenAIImagesRequestError(http.StatusBadRequest, err.Error()) } upstreamModel := account.GetMappedModel(requestModel) if err := validateOpenAIImagesModel(upstreamModel); err != nil { - return nil, err + return nil, newOpenAIImagesRequestError(http.StatusBadRequest, err.Error()) + } + if err := validateOpenAIImagesOptionsForModel(parsed, upstreamModel); err != nil { + return nil, newOpenAIImagesRequestError(http.StatusBadRequest, err.Error()) } logger.LegacyPrintf( "service.openai_gateway", @@ -545,6 +588,11 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesAPIKey( } upstreamCtx, releaseUpstreamCtx := s.detachOpenAIUpstreamContext(ctx) defer releaseUpstreamCtx() + if !parsed.Stream { + var cancelTotalTimeout context.CancelFunc + upstreamCtx, cancelTotalTimeout = s.openAIImagesNonstreamTotalContext(upstreamCtx) + defer cancelTotalTimeout() + } token, _, err := s.GetAccessToken(upstreamCtx, account) if err != nil { @@ -575,7 +623,10 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesAPIKey( Kind: "request_error", Message: safeErr, }) - return nil, fmt.Errorf("upstream request failed: %s", safeErr) + if ctx != nil && ctx.Err() != nil { + return nil, ctx.Err() + } + return nil, newOpenAIImagesStreamFailoverError(nil, http.StatusBadGateway, safeErr, false) } if resp.StatusCode >= 400 { respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) @@ -601,6 +652,9 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesAPIKey( RetryableOnSameAccount: shouldRetryOpenAIOnSamePoolAccount(account, resp.StatusCode, upstreamMsg, respBody), } } + if resp.StatusCode == http.StatusBadRequest { + return nil, newOpenAIImagesStreamFailoverError(resp, resp.StatusCode, upstreamMsg, false) + } return s.handleErrorResponse(upstreamCtx, resp, c, account, forwardBody, requestModel) } defer func() { _ = resp.Body.Close() }() @@ -619,7 +673,10 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesAPIKey( } else { nonStreamUsage, nonStreamCount, err := s.handleOpenAIImagesNonStreamingResponse(upstreamCtx, resp, c) if err != nil { - return nil, err + if ctx != nil && ctx.Err() != nil { + return nil, ctx.Err() + } + return nil, newOpenAIImagesStreamFailoverError(resp, http.StatusBadGateway, err.Error(), false) } usage = nonStreamUsage if nonStreamCount > 0 { @@ -640,6 +697,16 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesAPIKey( }, nil } +func (s *OpenAIGatewayService) openAIImagesNonstreamTotalContext(ctx context.Context) (context.Context, context.CancelFunc) { + if ctx == nil { + ctx = context.Background() + } + if s == nil || s.cfg == nil || s.cfg.Gateway.ImageNonstreamTotalTimeoutSeconds <= 0 { + return ctx, func() {} + } + return context.WithTimeout(ctx, time.Duration(s.cfg.Gateway.ImageNonstreamTotalTimeoutSeconds)*time.Second) +} + func (s *OpenAIGatewayService) buildOpenAIImagesRequest( ctx context.Context, c *gin.Context, @@ -785,9 +852,7 @@ func cloneMultipartHeader(src textproto.MIMEHeader) textproto.MIMEHeader { } func (s *OpenAIGatewayService) handleOpenAIImagesNonStreamingResponse(ctx context.Context, resp *http.Response, c *gin.Context) (OpenAIUsage, int, error) { - readCtx, cancelRead := s.detachedNonStreamingReadContext(ctx) - defer cancelRead() - body, err := ReadUpstreamResponseBodyWithContext(readCtx, resp.Body, s.cfg, c, openAITooLargeError) + body, err := ReadUpstreamResponseBodyWithContext(ctx, resp.Body, s.cfg, c, openAITooLargeError) if err != nil { return OpenAIUsage{}, 0, err } @@ -804,18 +869,378 @@ func (s *OpenAIGatewayService) handleOpenAIImagesNonStreamingResponse(ctx contex return usage, extractOpenAIImageCountFromJSONBytes(body), nil } +var errOpenAIImagesStreamIdleTimeout = errors.New("image stream data interval timeout") + +type openAIImagesStreamReadEvent struct { + line []byte + err error +} + +type openAIImagesStreamPump struct { + service *OpenAIGatewayService + ctx context.Context + resp *http.Response + c *gin.Context + flusher http.Flusher + events <-chan openAIImagesStreamReadEvent + stopReader chan struct{} + idleTimer *time.Timer + idleCh <-chan time.Time + idleTimeout time.Duration + keepaliveTicker *time.Ticker + keepaliveCh <-chan time.Time + keepaliveInterval time.Duration + clientDone <-chan struct{} + clientDisconnected bool + semanticOutput bool + nonSemanticOutput bool + lastDownstreamWrite time.Time + cancelDrainDeadline context.CancelFunc + disconnectLogMessage string + beforeSemanticWrite func() + semanticHeadersSet bool +} + +func newOpenAIImagesStreamPump( + s *OpenAIGatewayService, + ctx context.Context, + resp *http.Response, + c *gin.Context, + flusher http.Flusher, + disconnectLogMessage string, + beforeSemanticWrite func(), +) *openAIImagesStreamPump { + if ctx == nil { + ctx = context.Background() + } + stopReader := make(chan struct{}) + events := make(chan openAIImagesStreamReadEvent, 1) + go func() { + defer close(events) + reader := bufio.NewReader(resp.Body) + send := func(event openAIImagesStreamReadEvent) bool { + select { + case events <- event: + return true + case <-stopReader: + return false + } + } + for { + line, err := reader.ReadBytes('\n') + if len(line) > 0 && !send(openAIImagesStreamReadEvent{line: line}) { + return + } + if err != nil { + if err != io.EOF { + _ = send(openAIImagesStreamReadEvent{err: err}) + } + return + } + } + }() + + clientDone := ctx.Done() + if c != nil && c.Request != nil && c.Request.Context() != nil { + clientDone = c.Request.Context().Done() + } + pump := &openAIImagesStreamPump{ + service: s, + ctx: ctx, + resp: resp, + c: c, + flusher: flusher, + events: events, + stopReader: stopReader, + clientDone: clientDone, + lastDownstreamWrite: time.Now(), + disconnectLogMessage: disconnectLogMessage, + beforeSemanticWrite: beforeSemanticWrite, + } + if s != nil && s.cfg != nil && s.cfg.Gateway.StreamDataIntervalTimeout > 0 { + pump.idleTimeout = time.Duration(s.cfg.Gateway.StreamDataIntervalTimeout) * time.Second + pump.idleTimer = time.NewTimer(pump.idleTimeout) + pump.idleCh = pump.idleTimer.C + } + if s != nil && s.cfg != nil && s.cfg.Gateway.StreamKeepaliveInterval > 0 { + pump.keepaliveInterval = time.Duration(s.cfg.Gateway.StreamKeepaliveInterval) * time.Second + pump.keepaliveTicker = time.NewTicker(pump.keepaliveInterval) + pump.keepaliveCh = pump.keepaliveTicker.C + } + return pump +} + +func (p *openAIImagesStreamPump) Close() { + if p == nil { + return + } + if p.idleTimer != nil { + p.idleTimer.Stop() + } + if p.keepaliveTicker != nil { + p.keepaliveTicker.Stop() + } + if p.cancelDrainDeadline != nil { + p.cancelDrainDeadline() + } + select { + case <-p.stopReader: + default: + close(p.stopReader) + } + if p.resp != nil && p.resp.Body != nil { + _ = p.resp.Body.Close() + } +} + +func (p *openAIImagesStreamPump) resetIdleTimer() { + if p == nil || p.idleTimer == nil { + return + } + if !p.idleTimer.Stop() { + select { + case <-p.idleTimer.C: + default: + } + } + p.idleTimer.Reset(p.idleTimeout) +} + +func (p *openAIImagesStreamPump) startClientDisconnectDrain() { + if p == nil || p.clientDisconnected { + return + } + p.clientDisconnected = true + p.clientDone = nil + if p.cancelDrainDeadline == nil { + p.cancelDrainDeadline = p.service.startDisconnectedStreamDrainDeadline(p.ctx, p.resp.Body, p.resp.Header.Get("x-request-id")) + } + if strings.Contains(p.disconnectLogMessage, "OAuth") { + p.service.legacyLogClientDisconnectDrainDecision(p.ctx, "[OpenAI images OAuth] Client disconnected during streaming, continuing to drain upstream for usage") + } else { + p.service.legacyLogClientDisconnectDrainDecision(p.ctx, "[OpenAI images] Client disconnected during streaming, continuing to drain upstream for usage") + } +} + +func (p *openAIImagesStreamPump) markDownstreamWrite(semantic bool) { + if p == nil { + return + } + p.lastDownstreamWrite = time.Now() + if semantic { + p.semanticOutput = true + } else { + p.nonSemanticOutput = true + } +} + +func (p *openAIImagesStreamPump) write(data []byte, semantic bool) bool { + if p == nil || p.clientDisconnected { + return false + } + if p.clientDone != nil { + select { + case <-p.clientDone: + p.startClientDisconnectDrain() + return false + default: + } + } + if semantic && !p.semanticHeadersSet { + if p.beforeSemanticWrite != nil { + p.beforeSemanticWrite() + } + p.semanticHeadersSet = true + } + if _, err := p.c.Writer.Write(data); err != nil { + p.startClientDisconnectDrain() + return false + } + p.flusher.Flush() + p.markDownstreamWrite(semantic) + return true +} + +func (p *openAIImagesStreamPump) Next() ([]byte, error) { + for { + select { + case event, ok := <-p.events: + if !ok { + return nil, io.EOF + } + if len(event.line) > 0 { + p.resetIdleTimer() + return event.line, nil + } + if event.err != nil { + return nil, event.err + } + case <-p.idleCh: + _ = p.resp.Body.Close() + return nil, errOpenAIImagesStreamIdleTimeout + case <-p.clientDone: + p.startClientDisconnectDrain() + case <-p.keepaliveCh: + if p.clientDisconnected || time.Since(p.lastDownstreamWrite) < p.keepaliveInterval { + continue + } + p.write([]byte(":\n\n"), false) + } + } +} + +func (p *openAIImagesStreamPump) ClientDisconnected() bool { + return p != nil && p.clientDisconnected +} + +func (p *openAIImagesStreamPump) SemanticOutputWritten() bool { + return p != nil && p.semanticOutput +} + +func (p *openAIImagesStreamPump) SafeToFailoverAfterWrite() bool { + return p != nil && p.nonSemanticOutput && !p.semanticOutput +} + +func parseOpenAIImagesSSEEvent(event []byte) (string, []byte, bool) { + var eventName string + dataLines := make([]string, 0, 1) + for _, rawLine := range bytes.Split(event, []byte("\n")) { + line := strings.TrimRight(string(rawLine), "\r") + if strings.HasPrefix(line, "event:") { + eventName = strings.TrimSpace(strings.TrimPrefix(line, "event:")) + continue + } + if data, ok := extractOpenAISSEDataLine(line); ok { + dataLines = append(dataLines, data) + } + } + if len(dataLines) == 0 { + return eventName, nil, false + } + return eventName, []byte(strings.Join(dataLines, "\n")), true +} + +func openAIImagesDirectStreamFailure(eventName string, payload []byte) (int, string, bool) { + eventType := strings.ToLower(strings.TrimSpace(eventName)) + if gjson.ValidBytes(payload) { + if value := strings.ToLower(strings.TrimSpace(gjson.GetBytes(payload, "type").String())); value != "" { + eventType = value + } + } + failed := eventType == "error" || strings.HasSuffix(eventType, ".error") || strings.HasSuffix(eventType, ".failed") + if !failed && gjson.ValidBytes(payload) { + errorValue := gjson.GetBytes(payload, "error") + failed = errorValue.Exists() && errorValue.Type != gjson.Null && + (errorValue.IsObject() || errorValue.Type == gjson.String) + } + if !failed { + return 0, "", false + } + status := 0 + message := "" + if gjson.ValidBytes(payload) { + for _, path := range []string{"error.status", "error.status_code", "status", "status_code"} { + if value := int(gjson.GetBytes(payload, path).Int()); value > 0 { + status = value + break + } + } + for _, path := range []string{"error.message", "message"} { + if value := strings.TrimSpace(gjson.GetBytes(payload, path).String()); value != "" { + message = value + break + } + } + } + if status < http.StatusBadRequest || status > 599 { + status = http.StatusBadGateway + } + if message == "" { + message = "OpenAI image generation failed" + } + return status, sanitizeUpstreamErrorMessage(message), true +} + +func newOpenAIImagesStreamFailoverError(resp *http.Response, status int, message string, safeAfterWrite bool) *UpstreamFailoverError { + if status < http.StatusBadRequest || status > 599 { + status = http.StatusBadGateway + } + message = strings.TrimSpace(message) + if message == "" { + message = "upstream image request failed" + } + failoverErr := &UpstreamFailoverError{ + StatusCode: status, + ResponseBody: openAIImagesFailoverBody(message), + SafeToFailoverAfterWrite: safeAfterWrite, + } + if resp != nil { + failoverErr.ResponseHeaders = resp.Header.Clone() + } + if status == http.StatusBadRequest { + if isOpenAIImagesAccountCapabilityFailure(message) { + failoverErr.Scope = GatewayFailureScopeAccount + } else { + failoverErr.Scope = GatewayFailureScopeRequest + failoverErr.NextAccountAction = NextAccountStop + failoverErr.ClientStatusCode = http.StatusBadRequest + failoverErr.ClientMessage = strings.TrimSpace(message) + } + } + return failoverErr +} + +func isOpenAIImagesAccountCapabilityFailure(message string) bool { + message = strings.ToLower(strings.TrimSpace(message)) + for _, marker := range []string{ + "unsupported image model", + "image generation tool is not available", + "image_generation tool is not available", + "image tool is not available", + "not eligible for image generation", + "does not have access to image generation", + "does not have access to model", + "do not have access to model", + "you do not have access to it", + "account does not support", + "account is not supported", + "model is not supported for this account", + } { + if strings.Contains(message, marker) { + return true + } + } + return false +} + +func newOpenAIImagesRequestError(status int, message string) *UpstreamFailoverError { + if status < http.StatusBadRequest || status >= http.StatusInternalServerError { + status = http.StatusBadRequest + } + message = strings.TrimSpace(message) + if message == "" { + message = "invalid image request" + } + return &UpstreamFailoverError{ + StatusCode: status, + ResponseBody: openAIImagesFailoverBody(message), + Scope: GatewayFailureScopeRequest, + NextAccountAction: NextAccountStop, + ClientStatusCode: status, + ClientMessage: message, + } +} + func (s *OpenAIGatewayService) handleOpenAIImagesStreamingResponse( ctx context.Context, resp *http.Response, c *gin.Context, startTime time.Time, ) (OpenAIUsage, int, *int, error) { - responseheaders.WriteFilteredHeaders(c.Writer.Header(), resp.Header, s.responseHeaderFilter) contentType := strings.TrimSpace(resp.Header.Get("Content-Type")) if contentType == "" { contentType = "text/event-stream" } - c.Status(resp.StatusCode) c.Header("Content-Type", contentType) flusher, ok := c.Writer.(http.Flusher) @@ -823,56 +1248,76 @@ func (s *OpenAIGatewayService) handleOpenAIImagesStreamingResponse( return OpenAIUsage{}, 0, nil, fmt.Errorf("streaming is not supported by response writer") } - reader := bufio.NewReader(resp.Body) usage := OpenAIUsage{} imageCount := 0 var firstTokenMs *int - clientDisconnected := false - var cancelDisconnectedDrain context.CancelFunc - defer func() { - if cancelDisconnectedDrain != nil { - cancelDisconnectedDrain() + pump := newOpenAIImagesStreamPump( + s, + ctx, + resp, + c, + flusher, + "[OpenAI images] Client disconnected during streaming, continuing to drain upstream for usage", + func() { responseheaders.WriteFilteredHeaders(c.Writer.Header(), resp.Header, s.responseHeaderFilter) }, + ) + defer pump.Close() + var eventBuffer []byte + + processEvent := func(event []byte) error { + eventName, data, hasData := parseOpenAIImagesSSEEvent(event) + meaningfulData := hasData && strings.TrimSpace(string(data)) != "" && strings.TrimSpace(string(data)) != "[DONE]" + if meaningfulData { + if status, message, failed := openAIImagesDirectStreamFailure(eventName, data); failed { + if !pump.SemanticOutputWritten() && !pump.ClientDisconnected() { + return newOpenAIImagesStreamFailoverError(resp, status, message, pump.SafeToFailoverAfterWrite()) + } + pump.write(event, false) + return fmt.Errorf("upstream image generation failed: %s", message) + } + mergeOpenAIUsage(&usage, data) + if count := extractOpenAIImageCountFromJSONBytes(data); count > imageCount { + imageCount = count + } } - }() - startDisconnectedDrain := func() { - if cancelDisconnectedDrain == nil { - cancelDisconnectedDrain = s.startDisconnectedStreamDrainDeadline(ctx, resp.Body, resp.Header.Get("x-request-id")) + written := pump.write(event, meaningfulData) + if written && meaningfulData && firstTokenMs == nil { + ms := int(time.Since(startTime).Milliseconds()) + firstTokenMs = &ms } + return nil } for { - line, err := reader.ReadBytes('\n') + line, err := pump.Next() if len(line) > 0 { - if firstTokenMs == nil { - ms := int(time.Since(startTime).Milliseconds()) - firstTokenMs = &ms - } - if !clientDisconnected { - if _, writeErr := c.Writer.Write(line); writeErr != nil { - clientDisconnected = true - startDisconnectedDrain() - s.legacyLogClientDisconnectDrainDecision(ctx, "[OpenAI images] Client disconnected during streaming, continuing to drain upstream for usage") - } else { - flusher.Flush() - } - } - - if data, ok := extractOpenAISSEDataLine(strings.TrimRight(string(line), "\r\n")); ok && data != "" && data != "[DONE]" { - dataBytes := []byte(data) - mergeOpenAIUsage(&usage, dataBytes) - if count := extractOpenAIImageCountFromJSONBytes(dataBytes); count > imageCount { - imageCount = count + eventBuffer = append(eventBuffer, line...) + if len(bytes.TrimRight(line, "\r\n")) == 0 { + if processErr := processEvent(eventBuffer); processErr != nil { + return usage, imageCount, firstTokenMs, processErr } + eventBuffer = eventBuffer[:0] } } if err == io.EOF { + if len(eventBuffer) > 0 { + if processErr := processEvent(eventBuffer); processErr != nil { + return usage, imageCount, firstTokenMs, processErr + } + } break } if err != nil { - return OpenAIUsage{}, 0, firstTokenMs, err + if pump.ClientDisconnected() || (ctx != nil && ctx.Err() != nil) { + return usage, imageCount, firstTokenMs, fmt.Errorf("stream usage incomplete after disconnect: %w", err) + } + if !pump.SemanticOutputWritten() { + return usage, imageCount, firstTokenMs, newOpenAIImagesStreamFailoverError(resp, http.StatusBadGateway, err.Error(), pump.SafeToFailoverAfterWrite()) + } + _ = pump.write(append([]byte("event: error\ndata: "), append(buildOpenAIImagesStreamErrorBody(err.Error()), []byte("\n\n")...)...), false) + return usage, imageCount, firstTokenMs, err } } - if clientDisconnected { + if pump.ClientDisconnected() { if streamErr := s.clientDisconnectIncompleteUsageError(ctx); streamErr != nil { return usage, imageCount, firstTokenMs, streamErr } diff --git a/backend/internal/service/openai_images_responses.go b/backend/internal/service/openai_images_responses.go index 55b639fce..69be6057a 100644 --- a/backend/internal/service/openai_images_responses.go +++ b/backend/internal/service/openai_images_responses.go @@ -225,6 +225,9 @@ func buildOpenAIImagesResponsesRequest(parsed *OpenAIImagesRequest, toolModel st if prompt == "" { return nil, fmt.Errorf("prompt is required") } + if parsed.N > 1 { + return nil, fmt.Errorf("n greater than 1 is not supported for OAuth image accounts") + } inputImages := make([]string, 0, len(parsed.InputImageURLs)+len(parsed.Uploads)) for _, imageURL := range parsed.InputImageURLs { @@ -262,9 +265,6 @@ func buildOpenAIImagesResponsesRequest(parsed *OpenAIImagesRequest, toolModel st tool := []byte(`{"type":"image_generation","action":"","model":""}`) tool, _ = sjson.SetBytes(tool, "action", action) tool, _ = sjson.SetBytes(tool, "model", strings.TrimSpace(toolModel)) - if shouldPassOpenAIImagesN(toolModel, parsed.N) { - tool, _ = sjson.SetBytes(tool, "n", parsed.N) - } for _, field := range []struct { path string @@ -305,13 +305,6 @@ func buildOpenAIImagesResponsesRequest(parsed *OpenAIImagesRequest, toolModel st return req, nil } -func shouldPassOpenAIImagesN(model string, n int) bool { - if n <= 1 { - return false - } - return !strings.EqualFold(strings.TrimSpace(model), "dall-e-3") -} - func extractOpenAIImagesFromResponsesCompleted(payload []byte) ([]openAIResponsesImageResult, int64, []byte, openAIResponsesImageResult, error) { if gjson.GetBytes(payload, "type").String() != "response.completed" { return nil, 0, nil, openAIResponsesImageResult{}, fmt.Errorf("unexpected event type") @@ -592,6 +585,97 @@ func (s *OpenAIGatewayService) writeOpenAIImagesStreamEvent(c *gin.Context, flus return nil } +func (s *OpenAIGatewayService) readOpenAIImagesOAuthNonStreamingSSE(ctx context.Context, resp *http.Response, c *gin.Context) ([]byte, error) { + if resp == nil || resp.Body == nil { + return nil, fmt.Errorf("upstream response body is missing") + } + type readEvent struct { + line []byte + err error + } + events := make(chan readEvent, 1) + done := make(chan struct{}) + go func() { + defer close(events) + reader := bufio.NewReader(resp.Body) + send := func(event readEvent) bool { + select { + case events <- event: + return true + case <-done: + return false + } + } + for { + line, err := reader.ReadBytes('\n') + if len(line) > 0 && !send(readEvent{line: line}) { + return + } + if err != nil { + if err != io.EOF { + _ = send(readEvent{err: err}) + } + return + } + } + }() + defer close(done) + + var idleTimer *time.Timer + var idleCh <-chan time.Time + idleTimeout := time.Duration(0) + if s != nil && s.cfg != nil && s.cfg.Gateway.StreamDataIntervalTimeout > 0 { + idleTimeout = time.Duration(s.cfg.Gateway.StreamDataIntervalTimeout) * time.Second + idleTimer = time.NewTimer(idleTimeout) + idleCh = idleTimer.C + defer idleTimer.Stop() + } + resetIdleTimer := func() { + if idleTimer == nil { + return + } + if !idleTimer.Stop() { + select { + case <-idleTimer.C: + default: + } + } + idleTimer.Reset(idleTimeout) + } + + maxBytes := resolveUpstreamResponseReadLimit(s.cfg) + body := make([]byte, 0, 64<<10) + for { + select { + case event, ok := <-events: + if !ok { + return body, nil + } + if event.err != nil { + return nil, event.err + } + if len(event.line) == 0 { + continue + } + resetIdleTimer() + if int64(len(body))+int64(len(event.line)) > maxBytes { + _ = resp.Body.Close() + setOpsUpstreamError(c, http.StatusBadGateway, "upstream response too large", "") + return nil, fmt.Errorf("%w: limit=%d", ErrUpstreamResponseBodyTooLarge, maxBytes) + } + body = append(body, event.line...) + case <-idleCh: + _ = resp.Body.Close() + setOpsUpstreamError(c, http.StatusBadGateway, "upstream response read idle timeout", "") + return nil, errOpenAIImagesStreamIdleTimeout + case <-ctx.Done(): + _ = resp.Body.Close() + setOpsUpstreamError(c, http.StatusBadGateway, "upstream response total timeout", "") + return nil, ctx.Err() + } + } +} + func (s *OpenAIGatewayService) handleOpenAIImagesOAuthNonStreamingResponse( ctx context.Context, resp *http.Response, @@ -599,17 +683,12 @@ func (s *OpenAIGatewayService) handleOpenAIImagesOAuthNonStreamingResponse( responseFormat string, fallbackModel string, ) (OpenAIUsage, int, []string, error) { - readCtx, cancelRead := s.detachedNonStreamingReadContext(ctx) - defer cancelRead() - body, err := ReadUpstreamResponseBodyWithContext(readCtx, resp.Body, s.cfg, c, openAITooLargeError) + body, err := s.readOpenAIImagesOAuthNonStreamingSSE(ctx, resp, c) if err != nil { return OpenAIUsage{}, 0, nil, err } if status, message, failed := openAIImagesResponsesFailure(body); failed { - return OpenAIUsage{}, 0, nil, &UpstreamFailoverError{ - StatusCode: status, - ResponseBody: openAIImagesFailoverBody(message), - } + return OpenAIUsage{}, 0, nil, newOpenAIImagesStreamFailoverError(resp, status, message, false) } var usage OpenAIUsage @@ -627,10 +706,7 @@ func (s *OpenAIGatewayService) handleOpenAIImagesOAuthNonStreamingResponse( return OpenAIUsage{}, 0, nil, err } if len(results) == 0 { - return OpenAIUsage{}, 0, nil, &UpstreamFailoverError{ - StatusCode: http.StatusBadGateway, - ResponseBody: openAIImagesFailoverBody("upstream did not return image output"), - } + return OpenAIUsage{}, 0, nil, newOpenAIImagesStreamFailoverError(resp, http.StatusBadGateway, "upstream did not return image output", false) } if strings.TrimSpace(firstMeta.Model) == "" { firstMeta.Model = strings.TrimSpace(fallbackModel) @@ -654,12 +730,9 @@ func (s *OpenAIGatewayService) handleOpenAIImagesOAuthStreamingResponse( streamPrefix string, fallbackModel string, ) (OpenAIUsage, int, []string, *int, error) { - responseheaders.WriteFilteredHeaders(c.Writer.Header(), resp.Header, s.responseHeaderFilter) c.Header("Content-Type", "text/event-stream") c.Header("Cache-Control", "no-cache") c.Header("Connection", "keep-alive") - c.Status(resp.StatusCode) - flusher, ok := c.Writer.(http.Flusher) if !ok { return OpenAIUsage{}, 0, nil, nil, fmt.Errorf("streaming is not supported by response writer") @@ -670,7 +743,6 @@ func (s *OpenAIGatewayService) handleOpenAIImagesOAuthStreamingResponse( format = "b64_json" } - reader := bufio.NewReader(resp.Body) usage := OpenAIUsage{} imageCount := 0 var imageOutputSizes []string @@ -680,50 +752,57 @@ func (s *OpenAIGatewayService) handleOpenAIImagesOAuthStreamingResponse( pendingSeen := make(map[string]struct{}) streamMeta := openAIResponsesImageResult{Model: strings.TrimSpace(fallbackModel)} var createdAt int64 - clientDisconnected := false - var cancelDisconnectedDrain context.CancelFunc - defer func() { - if cancelDisconnectedDrain != nil { - cancelDisconnectedDrain() - } - }() - startDisconnectedDrain := func() { - if cancelDisconnectedDrain == nil { - cancelDisconnectedDrain = s.startDisconnectedStreamDrainDeadline(ctx, resp.Body, resp.Header.Get("x-request-id")) - } - } + pump := newOpenAIImagesStreamPump( + s, + ctx, + resp, + c, + flusher, + "[OpenAI images OAuth] Client disconnected during streaming, continuing to drain upstream for usage", + func() { responseheaders.WriteFilteredHeaders(c.Writer.Header(), resp.Header, s.responseHeaderFilter) }, + ) + defer pump.Close() emitEvent := func(eventName string, payload []byte) error { - if clientDisconnected { + if pump.ClientDisconnected() { return nil } - if writeErr := s.writeOpenAIImagesStreamEvent(c, flusher, eventName, payload); writeErr != nil { - clientDisconnected = true - startDisconnectedDrain() - s.legacyLogClientDisconnectDrainDecision(ctx, "[OpenAI images OAuth] Client disconnected during streaming, continuing to drain upstream for usage") - return nil + var event bytes.Buffer + if strings.TrimSpace(eventName) != "" { + _, _ = fmt.Fprintf(&event, "event: %s\n", eventName) + } + _, _ = fmt.Fprintf(&event, "data: %s\n\n", payload) + semantic := !strings.EqualFold(strings.TrimSpace(eventName), "error") + if pump.write(event.Bytes(), semantic) && semantic && firstTokenMs == nil { + ms := int(time.Since(startTime).Milliseconds()) + firstTokenMs = &ms } return nil } + streamFailure := func(status int, message string, cause error) error { + if cause == nil { + cause = fmt.Errorf("upstream image generation failed: %s", message) + } + if pump.ClientDisconnected() || (ctx != nil && ctx.Err() != nil) { + return cause + } + if !pump.SemanticOutputWritten() { + return newOpenAIImagesStreamFailoverError(resp, status, message, pump.SafeToFailoverAfterWrite()) + } + _ = emitEvent("error", buildOpenAIImagesStreamErrorBody(message)) + return cause + } for { - line, err := reader.ReadBytes('\n') + line, err := pump.Next() if len(line) > 0 { trimmedLine := strings.TrimRight(string(line), "\r\n") data, ok := extractOpenAISSEDataLine(trimmedLine) if ok && data != "" && data != "[DONE]" { - if firstTokenMs == nil { - ms := int(time.Since(startTime).Milliseconds()) - firstTokenMs = &ms - } dataBytes := []byte(data) s.parseSSEUsageBytes(dataBytes, &usage) if gjson.ValidBytes(dataBytes) { if status, message, failed := openAIImagesResponsesEventFailure(dataBytes); failed { - return usage, imageCount, imageOutputSizes, firstTokenMs, &UpstreamFailoverError{ - StatusCode: status, - ResponseBody: openAIImagesFailoverBody(message), - ResponseHeaders: resp.Header.Clone(), - } + return usage, imageCount, imageOutputSizes, firstTokenMs, streamFailure(status, message, nil) } if meta, eventCreatedAt, ok := extractOpenAIResponsesImageMetaFromLifecycleEvent(dataBytes); ok { mergeOpenAIResponsesImageMeta(&streamMeta, meta) @@ -754,8 +833,7 @@ func (s *OpenAIGatewayService) handleOpenAIImagesOAuthStreamingResponse( case "response.output_item.done": img, itemID, ok, extractErr := extractOpenAIImageFromResponsesOutputItemDone(dataBytes) if extractErr != nil { - _ = emitEvent("error", buildOpenAIImagesStreamErrorBody(extractErr.Error())) - return OpenAIUsage{}, imageCount, imageOutputSizes, firstTokenMs, extractErr + return usage, imageCount, imageOutputSizes, firstTokenMs, streamFailure(http.StatusBadGateway, extractErr.Error(), extractErr) } if !ok { break @@ -774,8 +852,7 @@ func (s *OpenAIGatewayService) handleOpenAIImagesOAuthStreamingResponse( case "response.completed": results, _, usageRaw, firstMeta, extractErr := extractOpenAIImagesFromResponsesCompleted(dataBytes) if extractErr != nil { - _ = emitEvent("error", buildOpenAIImagesStreamErrorBody(extractErr.Error())) - return OpenAIUsage{}, imageCount, imageOutputSizes, firstTokenMs, extractErr + return usage, imageCount, imageOutputSizes, firstTokenMs, streamFailure(http.StatusBadGateway, extractErr.Error(), extractErr) } mergeOpenAIResponsesImageMeta(&streamMeta, firstMeta) finalResults := make([]openAIResponsesImageResult, 0, len(results)+len(pendingResults)) @@ -791,8 +868,7 @@ func (s *OpenAIGatewayService) handleOpenAIImagesOAuthStreamingResponse( reconcileOpenAIResponsesImageResultSizes(finalResults, nil) if len(finalResults) == 0 { err = fmt.Errorf("upstream did not return image output") - _ = emitEvent("error", buildOpenAIImagesStreamErrorBody(err.Error())) - return OpenAIUsage{}, imageCount, imageOutputSizes, firstTokenMs, err + return usage, imageCount, imageOutputSizes, firstTokenMs, streamFailure(http.StatusBadGateway, err.Error(), err) } eventName := streamPrefix + ".completed" for _, img := range finalResults { @@ -806,7 +882,7 @@ func (s *OpenAIGatewayService) handleOpenAIImagesOAuthStreamingResponse( } imageCount = len(emitted) imageOutputSizes = openAIResponsesImageResultSizes(finalResults) - if clientDisconnected { + if pump.ClientDisconnected() { if streamErr := s.clientDisconnectIncompleteUsageError(ctx); streamErr != nil { return usage, imageCount, imageOutputSizes, firstTokenMs, streamErr } @@ -823,13 +899,12 @@ func (s *OpenAIGatewayService) handleOpenAIImagesOAuthStreamingResponse( break } if err != nil { - _ = emitEvent("error", buildOpenAIImagesStreamErrorBody(err.Error())) - return OpenAIUsage{}, imageCount, imageOutputSizes, firstTokenMs, err + return usage, imageCount, imageOutputSizes, firstTokenMs, streamFailure(http.StatusBadGateway, err.Error(), err) } } if imageCount > 0 { - if clientDisconnected { + if pump.ClientDisconnected() { if streamErr := s.clientDisconnectIncompleteUsageError(ctx); streamErr != nil { return usage, imageCount, imageOutputSizes, firstTokenMs, streamErr } @@ -857,7 +932,7 @@ func (s *OpenAIGatewayService) handleOpenAIImagesOAuthStreamingResponse( } imageCount = len(emitted) imageOutputSizes = openAIResponsesImageResultSizes(finalResults) - if clientDisconnected { + if pump.ClientDisconnected() { if streamErr := s.clientDisconnectIncompleteUsageError(ctx); streamErr != nil { return usage, imageCount, imageOutputSizes, firstTokenMs, streamErr } @@ -869,8 +944,7 @@ func (s *OpenAIGatewayService) handleOpenAIImagesOAuthStreamingResponse( } streamErr := fmt.Errorf("stream disconnected before image generation completed") - _ = emitEvent("error", buildOpenAIImagesStreamErrorBody(streamErr.Error())) - return OpenAIUsage{}, imageCount, imageOutputSizes, firstTokenMs, streamErr + return usage, imageCount, imageOutputSizes, firstTokenMs, streamFailure(http.StatusBadGateway, streamErr.Error(), streamErr) } func (s *OpenAIGatewayService) forwardOpenAIImagesOAuth( @@ -886,11 +960,17 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesOAuth( requestModel = mapped } if err := validateOpenAIImagesModel(requestModel); err != nil { - return nil, err + return nil, newOpenAIImagesRequestError(http.StatusBadRequest, err.Error()) } upstreamModel := account.GetMappedModel(requestModel) if err := validateOpenAIImagesModel(upstreamModel); err != nil { - return nil, err + return nil, newOpenAIImagesRequestError(http.StatusBadRequest, err.Error()) + } + if err := validateOpenAIImagesOptionsForModel(parsed, upstreamModel); err != nil { + return nil, newOpenAIImagesRequestError(http.StatusBadRequest, err.Error()) + } + if parsed.N > 1 { + return nil, newOpenAIImagesRequestError(http.StatusBadRequest, "n greater than 1 is not supported for OAuth image accounts") } logger.LegacyPrintf( "service.openai_gateway", @@ -903,6 +983,11 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesOAuth( ) upstreamCtx, releaseUpstreamCtx := s.detachOpenAIUpstreamContext(ctx) defer releaseUpstreamCtx() + if !parsed.Stream { + var cancelTotalTimeout context.CancelFunc + upstreamCtx, cancelTotalTimeout = s.openAIImagesNonstreamTotalContext(upstreamCtx) + defer cancelTotalTimeout() + } token, _, err := s.GetAccessToken(upstreamCtx, account) if err != nil { @@ -911,7 +996,7 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesOAuth( responsesBody, err := buildOpenAIImagesResponsesRequest(parsed, upstreamModel) if err != nil { - return nil, err + return nil, newOpenAIImagesRequestError(http.StatusBadRequest, err.Error()) } upstreamReq, err := s.buildUpstreamRequest(upstreamCtx, c, account, responsesBody, token, true, parsed.StickySessionSeed(), false) if err != nil { @@ -941,7 +1026,10 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesOAuth( Kind: "request_error", Message: safeErr, }) - return nil, fmt.Errorf("upstream request failed: %s", safeErr) + if ctx != nil && ctx.Err() != nil { + return nil, ctx.Err() + } + return nil, newOpenAIImagesStreamFailoverError(nil, http.StatusBadGateway, safeErr, false) } if resp.StatusCode >= 400 { respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) @@ -976,6 +1064,9 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesOAuth( RetryableOnSameAccount: shouldRetryOpenAIOnSamePoolAccount(account, resp.StatusCode, upstreamMsg, respBody), } } + if resp.StatusCode == http.StatusBadRequest { + return nil, newOpenAIImagesStreamFailoverError(resp, resp.StatusCode, upstreamMsg, false) + } return s.handleErrorResponse(upstreamCtx, resp, c, account, responsesBody, requestModel) } defer func() { _ = resp.Body.Close() }() @@ -994,7 +1085,14 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesOAuth( } else { usage, imageCount, imageOutputSizes, err = s.handleOpenAIImagesOAuthNonStreamingResponse(upstreamCtx, resp, c, parsed.ResponseFormat, requestModel) if err != nil { - return nil, err + if ctx != nil && ctx.Err() != nil { + return nil, ctx.Err() + } + var failoverErr *UpstreamFailoverError + if errors.As(err, &failoverErr) { + return nil, failoverErr + } + return nil, newOpenAIImagesStreamFailoverError(resp, http.StatusBadGateway, err.Error(), false) } } if imageCount <= 0 { diff --git a/backend/internal/service/openai_images_test.go b/backend/internal/service/openai_images_test.go index 4409ab1b4..918c85cd2 100644 --- a/backend/internal/service/openai_images_test.go +++ b/backend/internal/service/openai_images_test.go @@ -3,12 +3,14 @@ package service import ( "bytes" "context" + "errors" "io" "mime/multipart" "net/http" "net/http/httptest" "net/textproto" "strings" + "sync" "testing" "time" @@ -42,6 +44,37 @@ func TestOpenAIGatewayServiceParseOpenAIImagesRequest_JSON(t *testing.T) { require.False(t, parsed.Multipart) } +func TestOpenAIGatewayServiceParseOpenAIImagesRequest_ValidatesImageOptions(t *testing.T) { + tests := []struct { + name string + body string + message string + }{ + {name: "n below range", body: `{"model":"gpt-image-2","prompt":"cat","n":0}`, message: "n must be between 1 and 10"}, + {name: "n above range", body: `{"model":"gpt-image-2","prompt":"cat","n":11}`, message: "n must be between 1 and 10"}, + {name: "n fractional", body: `{"model":"gpt-image-2","prompt":"cat","n":1.5}`, message: "n must be an integer"}, + {name: "partial images below range", body: `{"model":"gpt-image-2","prompt":"cat","partial_images":-1}`, message: "partial_images must be between 0 and 3"}, + {name: "partial images above range", body: `{"model":"gpt-image-2","prompt":"cat","partial_images":4}`, message: "partial_images must be between 0 and 3"}, + {name: "compression below range", body: `{"model":"gpt-image-2","prompt":"cat","output_compression":-1}`, message: "output_compression must be between 0 and 100"}, + {name: "compression above range", body: `{"model":"gpt-image-2","prompt":"cat","output_compression":101}`, message: "output_compression must be between 0 and 100"}, + {name: "transparent gpt image 2", body: `{"model":"gpt-image-2","prompt":"cat","background":"transparent"}`, message: "background transparent is not supported by gpt-image-2"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + body := []byte(tt.body) + req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = req + + parsed, err := (&OpenAIGatewayService{}).ParseOpenAIImagesRequest(c, body) + require.Nil(t, parsed) + require.ErrorContains(t, err, tt.message) + }) + } +} + func TestOpenAIGatewayServiceParseOpenAIImagesRequest_MultipartEdit(t *testing.T) { gin.SetMode(gin.TestMode) @@ -76,6 +109,28 @@ func TestOpenAIGatewayServiceParseOpenAIImagesRequest_MultipartEdit(t *testing.T require.Equal(t, OpenAIImagesCapabilityNative, parsed.RequiredCapability) } +func TestOpenAIGatewayServiceParseOpenAIImagesRequest_RejectsOversizedMultipartPart(t *testing.T) { + gin.SetMode(gin.TestMode) + var body bytes.Buffer + writer := multipart.NewWriter(&body) + require.NoError(t, writer.WriteField("model", "gpt-image-2")) + require.NoError(t, writer.WriteField("prompt", "edit image")) + part, err := writer.CreateFormFile("image", "oversized.png") + require.NoError(t, err) + _, err = part.Write(bytes.Repeat([]byte{'x'}, openAIImageMaxUploadPartSize+1)) + require.NoError(t, err) + require.NoError(t, writer.Close()) + + req := httptest.NewRequest(http.MethodPost, "/v1/images/edits", bytes.NewReader(body.Bytes())) + req.Header.Set("Content-Type", writer.FormDataContentType()) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = req + + parsed, err := (&OpenAIGatewayService{}).ParseOpenAIImagesRequest(c, body.Bytes()) + require.Nil(t, parsed) + require.ErrorContains(t, err, "multipart field image exceeds the 20MB per-part limit") +} + func TestOpenAIGatewayServiceParseOpenAIImagesRequest_MultipartEditWithMaskAndNativeOptions(t *testing.T) { gin.SetMode(gin.TestMode) @@ -272,14 +327,20 @@ func TestResolveOpenAIImageBytes_PrefersInlineBase64(t *testing.T) { require.Equal(t, []byte("ABC"), data) } -func TestAccountSupportsOpenAIImageCapability_OAuthSupportsNative(t *testing.T) { - account := &Account{ +func TestAccountSupportsOpenAIImageCapability_PrefersNativeAPIKey(t *testing.T) { + oauthAccount := &Account{ Platform: PlatformOpenAI, Type: AccountTypeOAuth, } + apiKeyAccount := &Account{ + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + } - require.True(t, account.SupportsOpenAIImageCapability(OpenAIImagesCapabilityBasic)) - require.True(t, account.SupportsOpenAIImageCapability(OpenAIImagesCapabilityNative)) + require.True(t, oauthAccount.SupportsOpenAIImageCapability(OpenAIImagesCapabilityBasic)) + require.False(t, oauthAccount.SupportsOpenAIImageCapability(OpenAIImagesCapabilityNative)) + require.True(t, apiKeyAccount.SupportsOpenAIImageCapability(OpenAIImagesCapabilityBasic)) + require.True(t, apiKeyAccount.SupportsOpenAIImageCapability(OpenAIImagesCapabilityNative)) } func TestBuildOpenAIImagesURL_HandlesVersionedBaseURL(t *testing.T) { @@ -341,7 +402,7 @@ func findOpenAIImageTestSSEEvent(events []openAIImageTestSSEEvent, name string) func TestOpenAIGatewayServiceForwardImages_OAuthUsesResponsesAPI(t *testing.T) { gin.SetMode(gin.TestMode) - body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat","size":"1024x1024","quality":"high","n":2}`) + body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat","size":"1024x1024","quality":"high"}`) req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") @@ -405,7 +466,7 @@ func TestOpenAIGatewayServiceForwardImages_OAuthUsesResponsesAPI(t *testing.T) { require.Equal(t, "gpt-image-2", gjson.GetBytes(upstream.lastBody, "tools.0.model").String()) require.Equal(t, "1024x1024", gjson.GetBytes(upstream.lastBody, "tools.0.size").String()) require.Equal(t, "high", gjson.GetBytes(upstream.lastBody, "tools.0.quality").String()) - require.Equal(t, int64(2), gjson.GetBytes(upstream.lastBody, "tools.0.n").Int()) + require.False(t, gjson.GetBytes(upstream.lastBody, "tools.0.n").Exists()) require.Equal(t, "draw a cat", gjson.GetBytes(upstream.lastBody, "input.0.content.0.text").String()) require.Equal(t, http.StatusOK, rec.Code) @@ -414,6 +475,64 @@ func TestOpenAIGatewayServiceForwardImages_OAuthUsesResponsesAPI(t *testing.T) { require.Equal(t, "draw a cat", gjson.Get(rec.Body.String(), "data.0.revised_prompt").String()) } +func TestOpenAIGatewayServiceForwardImages_OAuthRejectsMultipleImagesBeforeUpstream(t *testing.T) { + gin.SetMode(gin.TestMode) + body := []byte(`{"model":"gpt-image-2","prompt":"draw cats","n":2}`) + req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = req + + upstream := &httpUpstreamRecorder{} + svc := &OpenAIGatewayService{httpUpstream: upstream} + parsed, err := svc.ParseOpenAIImagesRequest(c, body) + require.NoError(t, err) + account := &Account{ + ID: 1, Platform: PlatformOpenAI, Type: AccountTypeOAuth, + Credentials: map[string]any{"access_token": "token-123"}, + } + + result, err := svc.ForwardImages(context.Background(), c, account, body, parsed, "") + var requestErr *UpstreamFailoverError + require.ErrorAs(t, err, &requestErr) + require.Nil(t, result) + require.Equal(t, GatewayFailureScopeRequest, requestErr.Scope) + require.Equal(t, NextAccountStop, requestErr.NextAccountAction) + require.Equal(t, http.StatusBadRequest, requestErr.ClientStatusCode) + require.Contains(t, requestErr.ClientMessage, "n greater than 1") + require.Nil(t, upstream.lastReq) +} + +func TestOpenAIGatewayServiceForwardImages_RejectsTransparentAfterModelMapping(t *testing.T) { + gin.SetMode(gin.TestMode) + body := []byte(`{"model":"gpt-image-1","prompt":"draw a cat","background":"transparent"}`) + req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = req + + upstream := &httpUpstreamRecorder{} + svc := &OpenAIGatewayService{httpUpstream: upstream} + parsed, err := svc.ParseOpenAIImagesRequest(c, body) + require.NoError(t, err) + account := &Account{ + ID: 2, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, + Credentials: map[string]any{ + "api_key": "test-key", + "model_mapping": map[string]any{"gpt-image-1": "gpt-image-2"}, + }, + } + + result, err := svc.ForwardImages(context.Background(), c, account, body, parsed, "") + var requestErr *UpstreamFailoverError + require.ErrorAs(t, err, &requestErr) + require.Nil(t, result) + require.Equal(t, GatewayFailureScopeRequest, requestErr.Scope) + require.Equal(t, NextAccountStop, requestErr.NextAccountAction) + require.Contains(t, requestErr.ClientMessage, "transparent") + require.Nil(t, upstream.lastReq) +} + func TestOpenAIGatewayServiceForwardImages_OAuthAppliesAccountModelMapping(t *testing.T) { gin.SetMode(gin.TestMode) body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat"}`) @@ -512,6 +631,37 @@ func TestOpenAIGatewayServiceForwardImages_APIKeyGenerationUsesConfiguredV1BaseU require.Equal(t, "aGVsbG8=", gjson.Get(rec.Body.String(), "data.0.b64_json").String()) } +func TestOpenAIGatewayServiceForwardImages_APIKeyHTTP400PreservesRequestError(t *testing.T) { + gin.SetMode(gin.TestMode) + body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat"}`) + req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = req + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusBadRequest, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"error":{"message":"invalid output format"}}`)), + }} + svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream} + parsed, err := svc.ParseOpenAIImagesRequest(c, body) + require.NoError(t, err) + account := &Account{ + ID: 6, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, + Credentials: map[string]any{"api_key": "test-api-key"}, + } + + result, err := svc.ForwardImages(context.Background(), c, account, body, parsed, "") + var requestErr *UpstreamFailoverError + require.ErrorAs(t, err, &requestErr) + require.Nil(t, result) + require.Equal(t, GatewayFailureScopeRequest, requestErr.Scope) + require.Equal(t, NextAccountStop, requestErr.NextAccountAction) + require.Equal(t, "invalid output format", requestErr.ClientMessage) + require.Empty(t, recorder.Body.String()) +} + func TestOpenAIGatewayServiceForwardImages_APIKeyEditUsesConfiguredV1BaseURL(t *testing.T) { gin.SetMode(gin.TestMode) @@ -811,7 +961,7 @@ func TestOpenAIGatewayServiceForwardImages_OAuthEditsStreamingTransformsEvents(t require.False(t, gjson.Get(completed.Data, "revised_prompt").Exists()) } -func TestBuildOpenAIImagesResponsesRequest_PassesMultipleImageCount(t *testing.T) { +func TestBuildOpenAIImagesResponsesRequest_RejectsMultipleImageCount(t *testing.T) { parsed := &OpenAIImagesRequest{ Endpoint: openAIImagesGenerationsEndpoint, Model: "gpt-image-2", @@ -820,11 +970,8 @@ func TestBuildOpenAIImagesResponsesRequest_PassesMultipleImageCount(t *testing.T } body, err := buildOpenAIImagesResponsesRequest(parsed, "gpt-image-2") - require.NoError(t, err) - require.NotNil(t, body) - require.Equal(t, int64(2), gjson.GetBytes(body, "tools.0.n").Int()) - require.Equal(t, "gpt-image-2", gjson.GetBytes(body, "tools.0.model").String()) - require.Equal(t, "draw a cat", gjson.GetBytes(body, "input.0.content.0.text").String()) + require.Nil(t, body) + require.ErrorContains(t, err, "n greater than 1 is not supported for OAuth image accounts") } func TestBuildOpenAIImagesResponsesRequest_StripsInputFidelity(t *testing.T) { @@ -1026,6 +1173,8 @@ func TestOpenAIGatewayServiceForwardImages_OAuthStreamingPreservesUpstreamFailur require.ErrorAs(t, err, &failoverErr) require.Nil(t, result) require.Equal(t, http.StatusBadRequest, failoverErr.StatusCode) + require.Equal(t, GatewayFailureScopeAccount, failoverErr.Scope) + require.NotEqual(t, NextAccountStop, failoverErr.NextAccountAction) require.Contains(t, string(failoverErr.ResponseBody), "Unsupported image model gpt-image-2") require.Empty(t, rec.Body.String()) } @@ -1084,3 +1233,185 @@ func TestOpenAIImagesOAuthStreamingClientDisconnectKeepsBillableImageResult(t *t require.Equal(t, 5, usage.InputTokens) require.Equal(t, 9, usage.OutputTokens) } + +func TestOpenAIImagesDirectStreamFailureClassification(t *testing.T) { + tests := []struct { + name string + payload string + wantStop bool + wantStatus int + wantCapability bool + }{ + { + name: "request error stops without account failover", + payload: "event: error\ndata: {\"type\":\"error\",\"error\":{\"status\":400,\"message\":\"Unknown parameter: output_compression\"}}\n\n", + wantStop: true, + wantStatus: http.StatusBadRequest, + }, + { + name: "account capability error remains failover eligible", + payload: "event: error\ndata: {\"type\":\"error\",\"error\":{\"status\":400,\"message\":\"Unsupported image model gpt-image-2\"}}\n\n", + wantStatus: http.StatusBadRequest, + wantCapability: true, + }, + { + name: "rate limit remains failover eligible", + payload: "event: error\ndata: {\"type\":\"error\",\"error\":{\"status\":429,\"message\":\"rate limited\"}}\n\n", + wantStatus: http.StatusTooManyRequests, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil) + resp := &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader(tt.payload)), + } + + _, _, _, err := (&OpenAIGatewayService{}).handleOpenAIImagesStreamingResponse(context.Background(), resp, c, time.Now()) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, tt.wantStatus, failoverErr.StatusCode) + require.Equal(t, tt.wantStop, failoverErr.NextAccountAction == NextAccountStop) + if tt.wantCapability { + require.Equal(t, GatewayFailureScopeAccount, failoverErr.Scope) + } + require.Equal(t, tt.wantCapability, isOpenAIImagesAccountCapabilityFailure(failoverErr.ClientMessage) || isOpenAIImagesAccountCapabilityFailure(string(failoverErr.ResponseBody))) + require.Empty(t, recorder.Body.String()) + require.False(t, c.Writer.Written(), "pre-output failure must not commit HTTP 200") + }) + } +} + +func TestOpenAIImagesDirectStreamErrorNullIsNotFailure(t *testing.T) { + status, message, failed := openAIImagesDirectStreamFailure("image_generation.in_progress", []byte(`{"type":"image_generation.in_progress","error":null}`)) + require.False(t, failed) + require.Zero(t, status) + require.Empty(t, message) +} + +func TestOpenAIImagesDirectStreamKeepaliveDoesNotBlockFailover(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil) + reader, writer := io.Pipe() + defer writer.Close() + go func() { + time.Sleep(1100 * time.Millisecond) + _, _ = io.WriteString(writer, "event: error\ndata: {\"type\":\"error\",\"error\":{\"status\":429,\"message\":\"rate limited\"}}\n\n") + _ = writer.Close() + }() + resp := &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: reader, + } + svc := &OpenAIGatewayService{cfg: &config.Config{Gateway: config.GatewayConfig{StreamKeepaliveInterval: 1}}} + + _, _, _, err := svc.handleOpenAIImagesStreamingResponse(context.Background(), resp, c, time.Now()) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.True(t, failoverErr.SafeToFailoverAfterWrite) + require.Equal(t, ":\n\n", recorder.Body.String()) +} + +func TestOpenAIImagesDirectStreamIdleTimeoutBeforeOutputCanFailover(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil) + reader, writer := io.Pipe() + defer writer.Close() + resp := &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: reader, + } + svc := &OpenAIGatewayService{cfg: &config.Config{Gateway: config.GatewayConfig{StreamDataIntervalTimeout: 1}}} + + _, _, _, err := svc.handleOpenAIImagesStreamingResponse(context.Background(), resp, c, time.Now()) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Contains(t, string(failoverErr.ResponseBody), "image stream data interval timeout") + require.Empty(t, recorder.Body.String()) +} + +func TestOpenAIImagesOAuthNonStreamingSSEUsesIdleTimeout(t *testing.T) { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + reader, writer := io.Pipe() + defer writer.Close() + resp := &http.Response{StatusCode: http.StatusOK, Header: http.Header{}, Body: reader} + svc := &OpenAIGatewayService{cfg: &config.Config{Gateway: config.GatewayConfig{StreamDataIntervalTimeout: 1}}} + + started := time.Now() + body, err := svc.readOpenAIImagesOAuthNonStreamingSSE(context.Background(), resp, c) + require.Nil(t, body) + require.ErrorIs(t, err, errOpenAIImagesStreamIdleTimeout) + require.GreaterOrEqual(t, time.Since(started), time.Second) +} + +func TestOpenAIImagesNonstreamTotalContextUsesDedicatedLimit(t *testing.T) { + svc := &OpenAIGatewayService{cfg: &config.Config{Gateway: config.GatewayConfig{ImageNonstreamTotalTimeoutSeconds: 2}}} + ctx, cancel := svc.openAIImagesNonstreamTotalContext(context.Background()) + defer cancel() + deadline, ok := ctx.Deadline() + require.True(t, ok) + require.WithinDuration(t, time.Now().Add(2*time.Second), deadline, 250*time.Millisecond) +} + +type blockingOpenAIImageReadCloser struct { + closed chan struct{} + once sync.Once +} + +func (r *blockingOpenAIImageReadCloser) Read(_ []byte) (int, error) { + <-r.closed + return 0, io.EOF +} + +func (r *blockingOpenAIImageReadCloser) Close() error { + r.once.Do(func() { close(r.closed) }) + return nil +} + +func TestOpenAIImagesStreamClientCancellationClosesUpstreamImmediately(t *testing.T) { + gin.SetMode(gin.TestMode) + requestCtx, cancelRequest := context.WithCancel(context.Background()) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil).WithContext(requestCtx) + body := &blockingOpenAIImageReadCloser{closed: make(chan struct{})} + resp := &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: body, + } + result := make(chan error, 1) + svc := &OpenAIGatewayService{settingService: newOpenAIDetachedDrainSettingServiceForTest(t, false)} + go func() { + _, _, _, err := svc.handleOpenAIImagesStreamingResponse(requestCtx, resp, c, time.Now()) + result <- err + }() + + cancelRequest() + select { + case <-body.closed: + case <-time.After(time.Second): + t.Fatal("client cancellation did not close upstream body") + } + select { + case err := <-result: + require.Error(t, err) + var failoverErr *UpstreamFailoverError + require.False(t, errors.As(err, &failoverErr), "client cancellation must not switch accounts") + case <-time.After(time.Second): + t.Fatal("stream handler did not return after client cancellation") + } +} diff --git a/backend/internal/service/openai_oauth_passthrough_test.go b/backend/internal/service/openai_oauth_passthrough_test.go index 93519de93..7410e9b43 100644 --- a/backend/internal/service/openai_oauth_passthrough_test.go +++ b/backend/internal/service/openai_oauth_passthrough_test.go @@ -330,6 +330,136 @@ func TestOpenAIGatewayService_OAuthPassthrough_StreamKeepsToolNameAndBodyNormali require.NotContains(t, body, "\"name\":\"edit\"") } +func TestOpenAIGatewayService_OAuthPassthrough_ImageOnlyModelNormalizesBeforeForward(t *testing.T) { + gin.SetMode(gin.TestMode) + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(nil)) + + upstreamSSE := strings.Join([]string{ + `data: {"type":"response.completed","response":{"id":"resp_image","model":"gpt-5.4-mini","output":[{"type":"image_generation_call","result":"aGVsbG8="}],"usage":{"input_tokens":1,"output_tokens":2,"output_tokens_details":{"image_tokens":2}}}}`, + "", + "data: [DONE]", + "", + }, "\n") + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader(upstreamSSE)), + }} + svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream} + account := &Account{ + ID: 124, + Name: "oauth-image-passthrough", + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Concurrency: 1, + Credentials: map[string]any{"access_token": "oauth-token", "chatgpt_account_id": "chatgpt-acc"}, + Extra: map[string]any{"openai_passthrough": true}, + Status: StatusActive, + Schedulable: true, + } + + result, err := svc.Forward(context.Background(), c, account, []byte(`{"model":"gpt-image-2","prompt":"draw a cat","stream":true,"size":"1024x1024","n":1}`)) + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, "gpt-image-2", result.Model) + require.Equal(t, openAIImagesResponsesMainModel, gjson.GetBytes(upstream.lastBody, "model").String()) + require.Equal(t, "draw a cat", gjson.GetBytes(upstream.lastBody, "input").String()) + require.Equal(t, "image_generation", gjson.GetBytes(upstream.lastBody, "tools.0.type").String()) + require.Equal(t, "gpt-image-2", gjson.GetBytes(upstream.lastBody, "tools.0.model").String()) + require.Equal(t, "1024x1024", gjson.GetBytes(upstream.lastBody, "tools.0.size").String()) + require.False(t, gjson.GetBytes(upstream.lastBody, "prompt").Exists()) + require.False(t, gjson.GetBytes(upstream.lastBody, "n").Exists()) +} + +func TestOpenAIGatewayService_OAuthPassthrough_ImageOnlyModelRejectsUnsupportedOptions(t *testing.T) { + tests := []struct { + name string + body string + param string + messagePart string + }{ + { + name: "transparent background", + body: `{"model":"gpt-image-2","prompt":"draw a cat","background":"transparent"}`, + param: "background", + messagePart: "background transparent is not supported by gpt-image-2", + }, + { + name: "multiple images", + body: `{"model":"gpt-image-2","prompt":"draw a cat","n":2}`, + param: "n", + messagePart: "supports one image per request", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + gin.SetMode(gin.TestMode) + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(nil)) + upstream := &httpUpstreamRecorder{} + svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream} + account := &Account{ + ID: 126, + Name: "oauth-image-invalid-options", + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Concurrency: 1, + Credentials: map[string]any{"access_token": "oauth-token", "chatgpt_account_id": "chatgpt-acc"}, + Extra: map[string]any{"openai_passthrough": true}, + Status: StatusActive, + Schedulable: true, + } + + result, err := svc.Forward(context.Background(), c, account, []byte(tt.body)) + require.Error(t, err) + require.Nil(t, result) + require.Equal(t, http.StatusBadRequest, rec.Code) + require.Equal(t, tt.param, gjson.Get(rec.Body.String(), "error.param").String()) + require.Contains(t, rec.Body.String(), tt.messagePart) + require.Nil(t, upstream.lastReq) + }) + } +} + +func TestOpenAIGatewayService_OAuthPassthrough_ImageOnlyModelRejectsCodexStripPolicy(t *testing.T) { + gin.SetMode(gin.TestMode) + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(nil)) + c.Request.Header.Set("User-Agent", "codex_cli_rs/0.124.0") + upstream := &httpUpstreamRecorder{} + svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream} + account := &Account{ + ID: 127, + Name: "oauth-image-codex-strip", + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Concurrency: 1, + Credentials: map[string]any{"access_token": "oauth-token", "chatgpt_account_id": "chatgpt-acc"}, + Extra: map[string]any{ + "openai_passthrough": true, + featureKeyCodexImageGenerationExplicitToolPolicy: codexImageGenerationExplicitToolPolicyStrip, + }, + Status: StatusActive, + Schedulable: true, + } + + result, err := svc.Forward(context.Background(), c, account, []byte(`{"model":"gpt-image-2","prompt":"draw a cat"}`)) + require.Error(t, err) + require.Nil(t, result) + require.Equal(t, http.StatusBadRequest, rec.Code) + require.Contains(t, rec.Body.String(), "disabled by this account's Codex image generation policy") + require.Contains(t, rec.Body.String(), "/v1/images/generations") + require.Nil(t, upstream.lastReq) +} + func TestOpenAIGatewayService_OAuthPassthrough_CodexAutoReviewPreservesUpstreamModel(t *testing.T) { gin.SetMode(gin.TestMode) @@ -533,6 +663,46 @@ func TestOpenAIGatewayService_OAuthPassthrough_DisabledUsesLegacyTransform(t *te require.Contains(t, string(upstream.lastBody), `"stream":true`) } +func TestOpenAIGatewayService_OAuthLegacy_ImageOnlyModelStillNormalizes(t *testing.T) { + gin.SetMode(gin.TestMode) + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(nil)) + + upstreamSSE := strings.Join([]string{ + `data: {"type":"response.completed","response":{"id":"resp_image_legacy","model":"gpt-5.4-mini","output":[{"type":"image_generation_call","result":"aGVsbG8="}],"usage":{"input_tokens":1,"output_tokens":2,"output_tokens_details":{"image_tokens":2}}}}`, + "", + "data: [DONE]", + "", + }, "\n") + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader(upstreamSSE)), + }} + svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream} + account := &Account{ + ID: 125, + Name: "oauth-image-legacy", + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Concurrency: 1, + Credentials: map[string]any{"access_token": "oauth-token", "chatgpt_account_id": "chatgpt-acc"}, + Extra: map[string]any{"openai_passthrough": false}, + Status: StatusActive, + Schedulable: true, + } + + result, err := svc.Forward(context.Background(), c, account, []byte(`{"model":"gpt-image-2","prompt":"draw a cat","stream":true}`)) + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, "gpt-image-2", result.Model) + require.Equal(t, openAIImagesResponsesMainModel, gjson.GetBytes(upstream.lastBody, "model").String()) + require.Equal(t, "image_generation", gjson.GetBytes(upstream.lastBody, "tools.0.type").String()) + require.Equal(t, "gpt-image-2", gjson.GetBytes(upstream.lastBody, "tools.0.model").String()) +} + func TestOpenAIGatewayService_OAuthLegacy_CompositeCodexUAUsesCodexOriginator(t *testing.T) { gin.SetMode(gin.TestMode) @@ -1364,6 +1534,65 @@ func TestOpenAIGatewayService_APIKeyPassthrough_PreservesBodyAndUsesResponsesEnd require.Empty(t, upstream.lastReq.Header.Get("X-Test")) } +func TestOpenAIGatewayService_APIKeyPassthrough_RejectsImageOnlyModel(t *testing.T) { + gin.SetMode(gin.TestMode) + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(nil)) + upstream := &httpUpstreamRecorder{} + svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream} + account := &Account{ + ID: 457, + Name: "apikey-image-passthrough", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{"api_key": "sk-api-key", "base_url": "https://api.openai.com"}, + Extra: map[string]any{"openai_passthrough": true}, + Status: StatusActive, + Schedulable: true, + } + + result, err := svc.Forward(context.Background(), c, account, []byte(`{"model":"gpt-image-2","prompt":"draw a cat","stream":false}`)) + require.Error(t, err) + require.Nil(t, result) + require.Equal(t, http.StatusBadRequest, rec.Code) + require.Contains(t, rec.Body.String(), "does not accept image-only model") + require.Contains(t, rec.Body.String(), "/v1/images/generations") + require.Contains(t, rec.Body.String(), "Responses-compatible text model") + require.Nil(t, upstream.lastReq) +} + +func TestOpenAIGatewayService_APIKeyChatFallback_RejectsImageOnlyModel(t *testing.T) { + gin.SetMode(gin.TestMode) + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(nil)) + upstream := &httpUpstreamRecorder{} + svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream} + account := &Account{ + ID: 458, + Name: "apikey-chat-only", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{"api_key": "sk-api-key", "base_url": "https://api.openai.com"}, + Extra: map[string]any{"openai_responses_supported": false}, + Status: StatusActive, + Schedulable: true, + } + + result, err := svc.Forward(context.Background(), c, account, []byte(`{"model":"gpt-image-2","prompt":"draw a cat","stream":false}`)) + require.Error(t, err) + require.Nil(t, result) + require.Equal(t, http.StatusBadRequest, rec.Code) + require.Contains(t, rec.Body.String(), "does not accept image-only model") + require.Contains(t, rec.Body.String(), "/v1/images/generations") + require.Nil(t, upstream.lastReq) +} + func TestOpenAIGatewayService_OAuthPassthrough_WarnOnTimeoutHeadersForStream(t *testing.T) { gin.SetMode(gin.TestMode) logSink, restore := captureStructuredLog(t) diff --git a/deploy/config.example.yaml b/deploy/config.example.yaml index f0c61afc8..a52585b19 100644 --- a/deploy/config.example.yaml +++ b/deploy/config.example.yaml @@ -152,14 +152,16 @@ gateway: # Timeout for waiting upstream response headers (seconds) # 等待上游响应头超时时间（秒） response_header_timeout: 600 - # OpenAI/Codex 上游响应头超时（秒，0=禁用；避免长排队被通用超时截断） - openai_response_header_timeout: 0 + # OpenAI/Codex 上游响应头超时（秒，0=禁用） + openai_response_header_timeout: 600 + # Images 非流式请求总超时（秒，0=禁用）；独立于普通流数据间隔超时 + image_nonstream_total_timeout_seconds: 1800 # Max request body size in bytes (default: 256MB) # 请求体最大字节数（默认 256MB） max_body_size: 268435456 - # Max bytes to read for non-stream upstream responses (default: 8MB) - # 非流式上游响应体读取上限（默认 8MB） - upstream_response_read_max_bytes: 8388608 + # Max bytes to read for non-stream upstream responses (default: 128MB) + # 非流式上游响应体读取上限（默认 128MB） + upstream_response_read_max_bytes: 134217728 # Max bytes to read for proxy probe responses (default: 1MB) # 代理探测响应体读取上限（默认 1MB） proxy_probe_response_read_max_bytes: 1048576 diff --git a/docs/site/content/docs/operations/changelog.mdx b/docs/site/content/docs/operations/changelog.mdx index 4e5bcc269..2f5f233fc 100644 --- a/docs/site/content/docs/operations/changelog.mdx +++ b/docs/site/content/docs/operations/changelog.mdx @@ -3,6 +3,16 @@ title: 更新日志 description: 按版本记录 Pixel API 的主要功能更新。 --- +## v1.2.7 + +- 修复 `gpt-image-2` 长耗时生成被通用 180 秒读取时限中断的问题，为 Images 非流式请求增加独立总超时，并为 OAuth 内部 SSE 使用增量空闲检测。 +- 完善 Images 流式心跳、客户端断开监听和有界上游排空；首个语义事件前不再提前提交响应状态，传输失败可安全切换账号。 +- 区分请求参数型与账号能力型错误：请求级 400 精确返回且不影响账号健康，图片能力异常可在当前分组和备用 API Key 分组间继续故障转移。 +- 增加 `n`、`partial_images`、`output_compression`、上传大小和映射后模型校验，并拒绝 `gpt-image-2` 不支持的透明背景。 +- 调整 Images 调度和 `/responses` 兼容逻辑，原生 Image API 优先使用 API Key；OAuth 图片桥接不再发送非法的 `tools[0].n`。 +- 将 OpenAI 上游响应头默认超时调整为 600 秒、Images 非流式总超时调整为 1800 秒，并统一非流式响应读取上限示例为 128 MB。 +- 延长图片账号连通性测试时限，避免瞬态网络或上游容量错误直接将账号永久标记为异常；本次发布不修改数据库结构和历史数据。 + ## v1.2.4 - 修复 API 密钥列表批量用量查询中 PostgreSQL 时间参数被推断为 `text`，导致“今日”和“近 30 天”用量统一显示错误的问题。 From c5afe6c0515ce810403a336eddd308455c54f28c Mon Sep 17 00:00:00 2001 From: kouzhenqi <826563886@qq.com> Date: Thu, 23 Jul 2026 10:41:32 +0800 Subject: [PATCH 002/122] =?UTF-8?q?release:=20=E5=8F=91=E5=B8=83=201.2.14?= =?UTF-8?q?=20=E5=B9=B6=E5=AE=8C=E5=96=84=E7=BD=91=E5=85=B3=E3=80=81?= =?UTF-8?q?=E8=B4=A6=E5=8F=B7=E4=B8=8E=E8=BF=90=E7=BB=B4=E8=83=BD=E5=8A=9B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 完善 OpenAI、Grok 与 Anthropic 网关协议转换、故障转移、超时及 WebSocket 转发 - 加强账号导入、共享模式、Agent Identity 唯一性与调度缓存一致性 - 增加 API 密钥有效期、图片输入 Token 计费及用量时区统计能力 - 优化备份恢复、优雅停机、在线迁移和运维日志清理流程 - 重构用户控制台表格、筛选器、密钥、用量与主题交互体验 - 新增迁移 215 至 217、运维清理服务配置及配套自动化测试 - 将运行版本更新为 1.2.14 - 验证后端 go test ./...、前端生产构建与文档站生产构建 --- .gitattributes | 3 + backend/cmd/server/VERSION | 2 +- backend/cmd/server/cleanup_runner_test.go | 262 ++++ backend/cmd/server/main.go | 200 ++- backend/cmd/server/main_lifecycle_test.go | 363 ++++++ backend/cmd/server/main_migrate_only_test.go | 170 +++ backend/cmd/server/main_shutdown_test.go | 191 +++ backend/cmd/server/shutdown.go | 288 +++++ backend/cmd/server/wire.go | 56 +- backend/cmd/server/wire_gen.go | 62 +- backend/cmd/server/wire_gen_test.go | 7 +- backend/internal/config/config.go | 219 +++- backend/internal/config/config_test.go | 154 +++ .../internal/handler/admin/account_handler.go | 1 + ...ccount_handler_public_share_update_test.go | 65 + .../internal/handler/admin/setting_handler.go | 40 +- backend/internal/handler/api_key_handler.go | 28 +- .../api_key_handler_create_expiration_test.go | 64 + backend/internal/handler/auth_handler.go | 8 +- .../handler/auth_oauth_pending_flow.go | 2 +- backend/internal/handler/dto/mappers.go | 2 + .../handler/dto/mappers_usage_test.go | 21 + backend/internal/handler/dto/settings.go | 22 +- backend/internal/handler/dto/types.go | 6 +- backend/internal/handler/failover_loop.go | 19 +- .../internal/handler/failover_loop_test.go | 28 +- backend/internal/handler/gateway_helper.go | 8 + backend/internal/handler/grok_media.go | 147 ++- .../handler/openai_account_share_mode.go | 11 + .../internal/handler/openai_alpha_search.go | 8 +- .../handler/openai_chat_completions.go | 18 +- .../handler/openai_gateway_handler.go | 319 ++++- .../handler/openai_gateway_handler_test.go | 81 ++ backend/internal/handler/openai_images.go | 11 +- backend/internal/handler/ops_error_logger.go | 32 +- backend/internal/handler/setting_handler.go | 13 +- backend/internal/handler/usage_handler.go | 180 ++- ...usage_handler_dashboard_time_range_test.go | 297 +++++ ...user_account_agent_identity_import_test.go | 76 ++ .../internal/handler/user_account_handler.go | 101 +- .../handler/user_account_public_share_test.go | 359 ++++++ backend/internal/handler/wire.go | 30 +- .../pkg/apicompat/anthropic_responses_test.go | 14 + .../anthropic_to_responses_response.go | 86 +- .../anthropic_to_responses_stream_test.go | 184 +++ backend/internal/pkg/apicompat/types.go | 18 + backend/internal/pkg/claude/constants.go | 7 +- backend/internal/pkg/ip/ip.go | 79 +- backend/internal/pkg/ip/ip_test.go | 52 +- backend/internal/pkg/sysutil/restart.go | 83 +- backend/internal/pkg/sysutil/restart_test.go | 185 +++ backend/internal/pkg/tlsfingerprint/dialer.go | 240 +++- .../pkg/tlsfingerprint/dialer_test.go | 224 ++++ backend/internal/pkg/xai/billing.go | 2 + backend/internal/repository/account_repo.go | 104 +- .../account_repo_agent_identity_test.go | 134 ++ .../account_repo_grok_managed_extra_test.go | 21 + .../account_repo_integration_test.go | 11 + ...ccount_repo_public_share_scheduler_test.go | 79 ++ backend/internal/repository/ent.go | 91 +- .../repository/ent_migrate_only_test.go | 32 + .../repository/github_release_service.go | 47 +- .../repository/github_release_service_test.go | 34 + backend/internal/repository/http_upstream.go | 79 +- .../repository/http_upstream_redirect_test.go | 40 + .../internal/repository/http_upstream_test.go | 92 ++ .../internal/repository/migrations_runner.go | 486 ++++++- .../repository/migrations_runner_notx_test.go | 368 +++++- .../internal/repository/scheduler_cache.go | 5 + .../repository/scheduler_cache_unit_test.go | 25 + .../internal/repository/usage_billing_repo.go | 4 +- backend/internal/repository/usage_log_repo.go | 153 ++- .../usage_log_repo_dashboard_timezone_test.go | 107 ++ .../usage_log_repo_integration_test.go | 4 +- .../usage_log_repo_request_type_test.go | 24 + .../internal/repository/withdrawal_repo.go | 6 +- .../repository/withdrawal_repo_test.go | 82 ++ backend/internal/server/api_contract_test.go | 12 +- backend/internal/server/http.go | 68 +- .../internal/server/http_client_ip_test.go | 100 ++ .../server/middleware/api_key_auth.go | 50 +- .../server/middleware/api_key_auth_test.go | 96 ++ .../internal/server/middleware/middleware.go | 13 + backend/internal/server/routes/gateway.go | 14 + .../internal/server/routes/gateway_test.go | 16 + backend/internal/service/account.go | 105 +- .../service/account_credential_import.go | 81 ++ .../service/account_credential_import_test.go | 128 ++ .../service/account_credential_safety.go | 24 + .../service/account_grok_managed_extra.go | 85 ++ .../account_grok_managed_extra_test.go | 226 ++++ .../account_grok_media_eligibility_test.go | 77 ++ backend/internal/service/account_service.go | 495 +++++++- ...count_service_owned_agent_identity_test.go | 742 +++++++++++ .../internal/service/account_share_mode.go | 4 +- .../internal/service/account_usage_service.go | 4 +- backend/internal/service/admin_service.go | 117 +- .../service/admin_service_bulk_update_test.go | 140 +++ backend/internal/service/api_key_service.go | 87 +- .../api_key_service_create_expiration_test.go | 118 ++ .../service/api_key_service_delete_test.go | 49 + .../service/api_key_service_length_test.go | 46 + backend/internal/service/backup_service.go | 740 +++++++++-- .../internal/service/backup_service_test.go | 916 +++++++++++++- backend/internal/service/billing_service.go | 10 +- .../service/billing_service_unified_test.go | 8 +- .../service/channel_monitor_checker.go | 34 +- .../channel_monitor_checker_body_test.go | 20 + .../chatcompletions_anthropic_bridge.go | 12 +- backend/internal/service/domain_constants.go | 1 + .../service/functional_module_switch_test.go | 34 +- .../service/gateway_multiplatform_test.go | 13 +- backend/internal/service/gateway_request.go | 25 + .../internal/service/gateway_request_test.go | 18 + backend/internal/service/gateway_service.go | 74 +- backend/internal/service/grok_media.go | 398 +++++- .../service/grok_media_content_test.go | 294 +++++ backend/internal/service/grok_media_test.go | 171 +++ .../internal/service/grok_quota_service.go | 44 +- backend/internal/service/grok_upstream_url.go | 6 + .../internal/service/http_upstream_profile.go | 16 + .../service/http_upstream_profile_test.go | 10 + .../service/image_generation_intent.go | 65 + .../image_generation_intent_explicit_test.go | 94 ++ .../service/openai_account_model_transient.go | 164 +++ .../openai_account_model_transient_test.go | 72 ++ .../service/openai_account_runtime_block.go | 102 ++ .../service/openai_account_scheduler.go | 114 +- .../service/openai_account_scheduler_test.go | 101 +- .../internal/service/openai_alpha_search.go | 14 +- .../service/openai_codex_models_service.go | 98 +- .../openai_codex_models_service_test.go | 104 +- .../service/openai_codex_transform.go | 34 +- .../service/openai_codex_transform_test.go | 40 +- .../service/openai_fast_policy_ws_test.go | 77 ++ .../service/openai_first_output_timeout.go | 179 ++- .../openai_first_output_timeout_test.go | 41 + .../openai_gateway_chat_completions.go | 16 +- .../openai_gateway_chat_completions_raw.go | 16 +- .../openai_gateway_chat_completions_test.go | 103 ++ .../service/openai_gateway_messages.go | 17 +- .../openai_gateway_record_usage_test.go | 19 +- .../openai_gateway_responses_chat_fallback.go | 16 +- .../service/openai_gateway_service.go | 911 ++++++++++++-- .../service/openai_gateway_service_test.go | 86 +- .../internal/service/openai_grok_selection.go | 8 + backend/internal/service/openai_images.go | 14 +- .../service/openai_images_responses.go | 14 +- .../openai_request_body_limit_failover.go | 43 + ...openai_request_body_limit_failover_test.go | 43 + .../service/openai_responses_lite_tools.go | 39 +- .../openai_responses_lite_tools_test.go | 37 + .../openai_responses_rejected_field_retry.go | 133 ++ ...nai_responses_rejected_field_retry_test.go | 160 +++ .../internal/service/openai_ws_client_read.go | 71 ++ .../internal/service/openai_ws_forwarder.go | 101 +- ...penai_ws_forwarder_ingress_session_test.go | 201 ++- .../openai_ws_forwarder_success_test.go | 69 +- .../openai_ws_ratelimit_signal_test.go | 24 +- .../service/openai_ws_v2/passthrough_relay.go | 18 + .../openai_ws_v2/passthrough_relay_test.go | 39 +- .../openai_ws_v2_passthrough_adapter.go | 165 ++- .../internal/service/ops_cleanup_service.go | 716 +++++++---- .../service/ops_cleanup_service_test.go | 504 +++++++- backend/internal/service/ops_service.go | 11 + backend/internal/service/ops_settings.go | 48 +- .../service/ops_settings_advanced_test.go | 73 ++ .../scheduler_snapshot_hydration_test.go | 89 +- .../service/scheduler_snapshot_service.go | 3 + backend/internal/service/setting_service.go | 52 +- backend/internal/service/settings_view.go | 10 +- .../internal/service/token_refresh_service.go | 92 +- .../service/token_refresh_service_test.go | 135 ++ backend/internal/service/update_service.go | 28 +- .../service/update_service_security_test.go | 15 + backend/internal/service/usage_log.go | 2 + backend/internal/service/usage_service.go | 8 +- backend/internal/service/wire.go | 11 +- backend/internal/service/withdrawal.go | 35 +- backend/internal/web/embed_on.go | 56 +- backend/internal/web/embed_test.go | 18 + .../215_ops_daily_partition_shadow.sql | 364 ++++++ .../216_usage_log_image_input_tokens.sql | 9 + ...penai_owned_agent_identity_unique_notx.sql | 75 ++ ...openai_owned_agent_identity_unique_test.go | 69 + .../ops_daily_partition_shadow_test.go | 62 + deploy/.env.example | 9 + deploy/config.example.yaml | 56 +- deploy/docker-compose.dev.yml | 8 +- deploy/docker-compose.local.yml | 10 +- deploy/docker-compose.yml | 16 +- deploy/pixel-retention-cleanup.service | 18 + deploy/pixel-retention-cleanup.sh | 101 ++ deploy/pixel-retention-cleanup.timer | 11 + deploy/rsyslog.logrotate | 18 + frontend/src/App.vue | 21 +- frontend/src/__tests__/themeContrast.spec.ts | 71 ++ frontend/src/api/__tests__/keys.spec.ts | 52 + frontend/src/api/accounts.ts | 2 + frontend/src/api/admin/settings.ts | 2 + frontend/src/api/admin/system.ts | 1 + frontend/src/api/keys.ts | 7 +- frontend/src/api/usage.ts | 17 +- .../account/CredentialImportModal.vue | 226 ++-- .../admin/account/ImportDataModal.vue | 10 +- .../src/components/admin/usage/UsageTable.vue | 29 +- .../admin/user/UserBalanceModal.vue | 2 +- frontend/src/components/common/BaseDialog.vue | 285 ++++- frontend/src/components/common/DataTable.vue | 232 +++- .../src/components/common/DateRangePicker.vue | 258 +++- frontend/src/components/common/GroupBadge.vue | 10 +- frontend/src/components/common/Select.vue | 390 +++++- .../src/components/common/VersionBadge.vue | 3 +- .../common/__tests__/BaseDialog.spec.ts | 192 +++ .../common/__tests__/DataTable.spec.ts | 162 +++ .../common/__tests__/DateRangePicker.spec.ts | 207 ++- .../common/__tests__/Select.spec.ts | 377 ++++++ .../__tests__/uiSkinPropagation.spec.ts | 54 + .../src/components/keys/EndpointPopover.vue | 346 ++++- frontend/src/components/keys/UseKeyModal.vue | 385 ++++-- .../keys/__tests__/EndpointPopover.spec.ts | 330 ++++- .../keys/__tests__/UseKeyModal.spec.ts | 221 +++- frontend/src/components/layout/AppHeader.vue | 38 +- frontend/src/components/layout/AppLayout.vue | 11 +- frontend/src/components/layout/AppSidebar.vue | 5 +- .../src/components/layout/TablePageLayout.vue | 106 +- .../layout/__tests__/AppHeader.spec.ts | 22 + .../layout/__tests__/TablePageLayout.spec.ts | 23 + .../components/user/ImportAccountsModal.vue | 161 ++- ...ImportAccountsModal.agent-identity.spec.ts | 270 ++++ .../dashboard/UserAccountSharingStats.vue | 114 +- .../user/dashboard/UserDashboardCharts.vue | 54 +- .../dashboard/UserDashboardQuickActions.vue | 44 +- .../dashboard/UserDashboardRecentUsage.vue | 32 +- .../user/dashboard/UserDashboardStats.vue | 178 +-- .../__tests__/UserAccountSharingStats.spec.ts | 94 ++ .../user/profile/ProfileWithdrawalCard.vue | 6 +- .../__tests__/ProfileWithdrawalCard.spec.ts | 4 +- .../composables/__tests__/useUiSkin.spec.ts | 27 + frontend/src/composables/useUiSkin.ts | 14 + frontend/src/i18n/locales/en.ts | 76 +- frontend/src/i18n/locales/zh.ts | 77 +- frontend/src/main.ts | 2 + frontend/src/router/index.ts | 3 + frontend/src/router/meta.d.ts | 5 + frontend/src/stores/app.ts | 5 + frontend/src/style.css | 784 +++++++++++- frontend/src/types/index.ts | 4 + .../utils/__tests__/apiKeyExpiration.spec.ts | 43 + frontend/src/utils/apiKeyExpiration.ts | 34 + frontend/src/utils/branding.ts | 19 + frontend/src/utils/imageUsage.ts | 16 + frontend/src/views/admin/AccountsView.vue | 38 +- frontend/src/views/admin/SettingsView.vue | 99 +- .../ops/components/OpsSettingsDialog.vue | 2 +- frontend/src/views/user/AccountsView.vue | 7 + frontend/src/views/user/DashboardView.vue | 139 +- frontend/src/views/user/KeysView.vue | 852 +++++++++---- frontend/src/views/user/ProfileView.vue | 4 + frontend/src/views/user/UsageView.vue | 1119 +++++++++-------- .../user/__tests__/DashboardView.spec.ts | 402 ++++++ .../src/views/user/__tests__/KeysView.spec.ts | 462 +++++++ .../views/user/__tests__/ProfileView.spec.ts | 8 +- .../views/user/__tests__/UsageView.spec.ts | 307 ++++- frontend/tailwind.config.js | 36 +- 265 files changed, 25726 insertions(+), 3223 deletions(-) create mode 100644 backend/cmd/server/cleanup_runner_test.go create mode 100644 backend/cmd/server/main_lifecycle_test.go create mode 100644 backend/cmd/server/main_migrate_only_test.go create mode 100644 backend/cmd/server/main_shutdown_test.go create mode 100644 backend/cmd/server/shutdown.go create mode 100644 backend/internal/handler/admin/account_handler_public_share_update_test.go create mode 100644 backend/internal/handler/api_key_handler_create_expiration_test.go create mode 100644 backend/internal/handler/usage_handler_dashboard_time_range_test.go create mode 100644 backend/internal/handler/user_account_agent_identity_import_test.go create mode 100644 backend/internal/pkg/apicompat/anthropic_to_responses_stream_test.go create mode 100644 backend/internal/pkg/sysutil/restart_test.go create mode 100644 backend/internal/repository/account_repo_agent_identity_test.go create mode 100644 backend/internal/repository/account_repo_grok_managed_extra_test.go create mode 100644 backend/internal/repository/account_repo_public_share_scheduler_test.go create mode 100644 backend/internal/repository/ent_migrate_only_test.go create mode 100644 backend/internal/repository/http_upstream_redirect_test.go create mode 100644 backend/internal/repository/usage_log_repo_dashboard_timezone_test.go create mode 100644 backend/internal/repository/withdrawal_repo_test.go create mode 100644 backend/internal/server/http_client_ip_test.go create mode 100644 backend/internal/service/account_grok_managed_extra.go create mode 100644 backend/internal/service/account_grok_managed_extra_test.go create mode 100644 backend/internal/service/account_grok_media_eligibility_test.go create mode 100644 backend/internal/service/account_service_owned_agent_identity_test.go create mode 100644 backend/internal/service/api_key_service_create_expiration_test.go create mode 100644 backend/internal/service/api_key_service_length_test.go create mode 100644 backend/internal/service/grok_media_content_test.go create mode 100644 backend/internal/service/image_generation_intent_explicit_test.go 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_request_body_limit_failover.go create mode 100644 backend/internal/service/openai_request_body_limit_failover_test.go 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 create mode 100644 backend/internal/service/openai_ws_client_read.go create mode 100644 backend/internal/service/update_service_security_test.go create mode 100644 backend/migrations/215_ops_daily_partition_shadow.sql create mode 100644 backend/migrations/216_usage_log_image_input_tokens.sql create mode 100644 backend/migrations/217_openai_owned_agent_identity_unique_notx.sql create mode 100644 backend/migrations/openai_owned_agent_identity_unique_test.go create mode 100644 backend/migrations/ops_daily_partition_shadow_test.go create mode 100644 deploy/pixel-retention-cleanup.service create mode 100644 deploy/pixel-retention-cleanup.sh create mode 100644 deploy/pixel-retention-cleanup.timer create mode 100644 deploy/rsyslog.logrotate create mode 100644 frontend/src/__tests__/themeContrast.spec.ts create mode 100644 frontend/src/api/__tests__/keys.spec.ts create mode 100644 frontend/src/components/common/__tests__/BaseDialog.spec.ts create mode 100644 frontend/src/components/common/__tests__/DataTable.spec.ts create mode 100644 frontend/src/components/common/__tests__/Select.spec.ts create mode 100644 frontend/src/components/common/__tests__/uiSkinPropagation.spec.ts create mode 100644 frontend/src/components/layout/__tests__/AppHeader.spec.ts create mode 100644 frontend/src/components/layout/__tests__/TablePageLayout.spec.ts create mode 100644 frontend/src/components/user/__tests__/ImportAccountsModal.agent-identity.spec.ts create mode 100644 frontend/src/components/user/dashboard/__tests__/UserAccountSharingStats.spec.ts create mode 100644 frontend/src/composables/__tests__/useUiSkin.spec.ts create mode 100644 frontend/src/composables/useUiSkin.ts create mode 100644 frontend/src/utils/__tests__/apiKeyExpiration.spec.ts create mode 100644 frontend/src/utils/apiKeyExpiration.ts create mode 100644 frontend/src/utils/branding.ts create mode 100644 frontend/src/utils/imageUsage.ts create mode 100644 frontend/src/views/user/__tests__/DashboardView.spec.ts create mode 100644 frontend/src/views/user/__tests__/KeysView.spec.ts diff --git a/.gitattributes b/.gitattributes index 37e3bee2c..aa687a76d 100644 --- a/.gitattributes +++ b/.gitattributes @@ -13,6 +13,9 @@ backend/migrations/*.sql text eol=lf # Shell 脚本 *.sh text eol=lf +*.service text eol=lf +*.timer text eol=lf +*.logrotate text eol=lf # YAML/YML 配置文件 *.yaml text eol=lf diff --git a/backend/cmd/server/VERSION b/backend/cmd/server/VERSION index c04c650a7..fd9d1a5ac 100644 --- a/backend/cmd/server/VERSION +++ b/backend/cmd/server/VERSION @@ -1 +1 @@ -1.2.7 +1.2.14 diff --git a/backend/cmd/server/cleanup_runner_test.go b/backend/cmd/server/cleanup_runner_test.go new file mode 100644 index 000000000..d63608ef9 --- /dev/null +++ b/backend/cmd/server/cleanup_runner_test.go @@ -0,0 +1,262 @@ +package main + +import ( + "context" + "errors" + "strings" + "sync" + "sync/atomic" + "testing" + "time" +) + +func TestRunProcessCleanupWaitsForParallelStepsBeforeSequentialInfra(t *testing.T) { + applicationStarted := make(chan struct{}, 2) + releaseApplication := make(chan struct{}) + var releaseOnce sync.Once + defer releaseOnce.Do(func() { close(releaseApplication) }) + + var applicationCompleted atomic.Int32 + parallelSteps := []cleanupStep{ + { + name: "application-a", + fn: func() error { + applicationStarted <- struct{}{} + <-releaseApplication + applicationCompleted.Add(1) + return nil + }, + }, + { + name: "application-b", + fn: func() error { + applicationStarted <- struct{}{} + <-releaseApplication + applicationCompleted.Add(1) + return nil + }, + }, + } + + infraStartedTooEarly := errors.New("infrastructure cleanup started before application cleanup completed") + infraOutOfOrder := errors.New("infrastructure cleanup ran out of order") + var firstInfraCompleted atomic.Bool + var secondInfraCompleted atomic.Bool + infraSteps := []cleanupStep{ + { + name: "redis", + fn: func() error { + if applicationCompleted.Load() != int32(len(parallelSteps)) { + return infraStartedTooEarly + } + firstInfraCompleted.Store(true) + return nil + }, + }, + { + name: "ent", + fn: func() error { + if applicationCompleted.Load() != int32(len(parallelSteps)) { + return infraStartedTooEarly + } + if !firstInfraCompleted.Load() { + return infraOutOfOrder + } + secondInfraCompleted.Store(true) + return nil + }, + }, + } + + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + result := make(chan error, 1) + go func() { + result <- runProcessCleanup(ctx, parallelSteps, infraSteps) + }() + + waitForCleanupSignal(t, applicationStarted, "first application cleanup step to start") + waitForCleanupSignal(t, applicationStarted, "second application cleanup step to start") + releaseOnce.Do(func() { close(releaseApplication) }) + + if err := waitForCleanupResult(t, result); err != nil { + t.Fatalf("runProcessCleanup() error = %v", err) + } + if got := applicationCompleted.Load(); got != int32(len(parallelSteps)) { + t.Fatalf("completed application steps = %d, want %d", got, len(parallelSteps)) + } + if !firstInfraCompleted.Load() { + t.Fatal("first infrastructure cleanup step did not complete") + } + if !secondInfraCompleted.Load() { + t.Fatal("second infrastructure cleanup step did not complete") + } +} + +func TestRunProcessCleanupApplicationTimeoutSkipsInfra(t *testing.T) { + applicationStarted := make(chan struct{}) + releaseApplication := make(chan struct{}) + applicationExited := make(chan struct{}) + infraStarted := make(chan struct{}, 1) + var releaseOnce sync.Once + defer releaseOnce.Do(func() { close(releaseApplication) }) + + parallelSteps := []cleanupStep{ + { + name: "blocked-application", + fn: func() error { + close(applicationStarted) + defer close(applicationExited) + <-releaseApplication + return nil + }, + }, + } + infraSteps := []cleanupStep{ + { + name: "redis", + fn: func() error { + infraStarted <- struct{}{} + return nil + }, + }, + } + + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + result := make(chan error, 1) + go func() { + result <- runProcessCleanup(ctx, parallelSteps, infraSteps) + }() + + waitForCleanupSignal(t, applicationStarted, "blocked application cleanup step to start") + err := waitForCleanupResult(t, result) + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("runProcessCleanup() error = %v, want wrapping %v", err, context.DeadlineExceeded) + } + if !strings.Contains(err.Error(), "blocked-application") { + t.Fatalf("runProcessCleanup() error = %v, want pending step name", err) + } + assertNoCleanupSignal(t, infraStarted, "infrastructure cleanup started after application timeout") + + releaseOnce.Do(func() { close(releaseApplication) }) + waitForCleanupSignal(t, applicationExited, "blocked application cleanup goroutine to exit") + assertNoCleanupSignal(t, infraStarted, "infrastructure cleanup started after timed-out application step was released") +} + +func TestRunProcessCleanupInfraTimeoutStopsRemainingInfra(t *testing.T) { + firstInfraStarted := make(chan struct{}) + releaseFirstInfra := make(chan struct{}) + firstInfraExited := make(chan struct{}) + secondInfraStarted := make(chan struct{}, 1) + var releaseOnce sync.Once + defer releaseOnce.Do(func() { close(releaseFirstInfra) }) + + infraSteps := []cleanupStep{ + { + name: "blocked-redis", + fn: func() error { + close(firstInfraStarted) + defer close(firstInfraExited) + <-releaseFirstInfra + return nil + }, + }, + { + name: "ent", + fn: func() error { + secondInfraStarted <- struct{}{} + return nil + }, + }, + } + + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + result := make(chan error, 1) + go func() { + result <- runProcessCleanup(ctx, nil, infraSteps) + }() + + waitForCleanupSignal(t, firstInfraStarted, "first infrastructure cleanup step to start") + err := waitForCleanupResult(t, result) + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("runProcessCleanup() error = %v, want wrapping %v", err, context.DeadlineExceeded) + } + if !strings.Contains(err.Error(), "blocked-redis") { + t.Fatalf("runProcessCleanup() error = %v, want pending step name", err) + } + assertNoCleanupSignal(t, secondInfraStarted, "second infrastructure cleanup step started after timeout") + + releaseOnce.Do(func() { close(releaseFirstInfra) }) + waitForCleanupSignal(t, firstInfraExited, "blocked infrastructure cleanup goroutine to exit") + assertNoCleanupSignal(t, secondInfraStarted, "second infrastructure cleanup step started after timed-out step was released") +} + +func TestRunProcessCleanupStepErrorsDoNotStopRemainingSteps(t *testing.T) { + applicationErr := errors.New("application cleanup failed") + infraErr := errors.New("infrastructure cleanup failed") + var successfulApplicationCalls atomic.Int32 + var successfulInfraCalls atomic.Int32 + + parallelSteps := []cleanupStep{ + {name: "failed-application", fn: func() error { return applicationErr }}, + {name: "successful-application", fn: func() error { + successfulApplicationCalls.Add(1) + return nil + }}, + } + infraSteps := []cleanupStep{ + {name: "failed-infra", fn: func() error { return infraErr }}, + {name: "successful-infra", fn: func() error { + successfulInfraCalls.Add(1) + return nil + }}, + } + + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + err := runProcessCleanup(ctx, parallelSteps, infraSteps) + + if !errors.Is(err, applicationErr) { + t.Errorf("runProcessCleanup() error = %v, want wrapping application error %v", err, applicationErr) + } + if !errors.Is(err, infraErr) { + t.Errorf("runProcessCleanup() error = %v, want wrapping infrastructure error %v", err, infraErr) + } + if got := successfulApplicationCalls.Load(); got != 1 { + t.Errorf("successful application cleanup calls = %d, want 1", got) + } + if got := successfulInfraCalls.Load(); got != 1 { + t.Errorf("successful infrastructure cleanup calls = %d, want 1", got) + } +} + +func waitForCleanupSignal(t *testing.T, signal <-chan struct{}, description string) { + t.Helper() + select { + case <-signal: + case <-time.After(time.Second): + t.Fatalf("timed out waiting for %s", description) + } +} + +func waitForCleanupResult(t *testing.T, result <-chan error) error { + t.Helper() + select { + case err := <-result: + return err + case <-time.After(time.Second): + t.Fatal("timed out waiting for runProcessCleanup to return") + return nil + } +} + +func assertNoCleanupSignal(t *testing.T, signal <-chan struct{}, failureMessage string) { + t.Helper() + select { + case <-signal: + t.Fatal(failureMessage) + case <-time.After(50 * time.Millisecond): + } +} diff --git a/backend/cmd/server/main.go b/backend/cmd/server/main.go index 96d46a3ae..4d6d2c026 100644 --- a/backend/cmd/server/main.go +++ b/backend/cmd/server/main.go @@ -24,6 +24,7 @@ import ( "github.com/Wei-Shaw/sub2api/internal/config" "github.com/Wei-Shaw/sub2api/internal/handler" "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/Wei-Shaw/sub2api/internal/repository" "github.com/Wei-Shaw/sub2api/internal/server/middleware" "github.com/Wei-Shaw/sub2api/internal/setup" "github.com/Wei-Shaw/sub2api/internal/web" @@ -42,6 +43,66 @@ var ( BuildType = "source" // "source" for manual builds, "release" for CI builds (set by ldflags) ) +const defaultMigrationTimeout = 30 * time.Minute + +type commandOptions struct { + setupMode bool + showVersion bool + migrateOnly bool + migrationTimeout time.Duration +} + +func (o commandOptions) validate() error { + if o.migrateOnly && o.setupMode { + return errors.New("--migrate-only cannot be combined with --setup") + } + if o.migrateOnly && o.showVersion { + return errors.New("--migrate-only cannot be combined with --version") + } + if o.migrateOnly && o.migrationTimeout <= 0 { + return errors.New("--migration-timeout must be greater than zero") + } + return nil +} + +type bootstrapConfigLoader func() (*config.Config, error) +type configuredMigrationRunner func(context.Context, *config.Config) error + +func runMigrationsOnly( + parent context.Context, + timeout time.Duration, + loadConfig bootstrapConfigLoader, + runMigrations configuredMigrationRunner, +) error { + if parent == nil { + return errors.New("nil parent context") + } + if timeout <= 0 { + return errors.New("migration timeout must be greater than zero") + } + if loadConfig == nil { + return errors.New("nil bootstrap config loader") + } + if runMigrations == nil { + return errors.New("nil configured migration runner") + } + + cfg, err := loadConfig() + if err != nil { + return fmt.Errorf("load migration config: %w", err) + } + if cfg == nil { + return errors.New("load migration config: nil config") + } + + ctx, cancel := context.WithTimeout(parent, timeout) + defer cancel() + if err := runMigrations(ctx, cfg); err != nil { + return fmt.Errorf("run database migrations: %w", err) + } + return nil +} + func init() { // 如果 Version 已通过 ldflags 注入（例如 -X main.Version=...），则不要覆盖。 if strings.TrimSpace(Version) != "" { @@ -64,13 +125,40 @@ func main() { // Parse command line flags setupMode := flag.Bool("setup", false, "Run setup wizard in CLI mode") showVersion := flag.Bool("version", false, "Show version information") + migrateOnly := flag.Bool("migrate-only", false, "Run embedded database migrations and exit") + migrationTimeout := flag.Duration("migration-timeout", defaultMigrationTimeout, "Maximum duration for --migrate-only") flag.Parse() + options := commandOptions{ + setupMode: *setupMode, + showVersion: *showVersion, + migrateOnly: *migrateOnly, + migrationTimeout: *migrationTimeout, + } + if err := options.validate(); err != nil { + log.Fatalf("Invalid command options: %v", err) + } + if *showVersion { log.Printf("Sub2API %s (commit: %s, built: %s)\n", Version, Commit, Date) return } + if *migrateOnly { + migrationParent, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) + defer stop() + if err := runMigrationsOnly( + migrationParent, + *migrationTimeout, + config.LoadForBootstrap, + repository.ApplyConfiguredMigrations, + ); err != nil { + log.Fatalf("Database migration failed: %v", err) + } + log.Println("Database migrations completed successfully") + return + } + // CLI setup mode if *setupMode { if err := setup.RunCLI(); err != nil { @@ -96,7 +184,15 @@ func main() { } // Normal server mode - runMainServer() + if err := runMainServer(); err != nil { + if errors.Is(err, errServerRestartRequested) { + log.Println("Graceful cleanup completed; exiting for process supervisor restart") + } else { + log.Printf("Server terminated: %v", err) + } + logger.Sync() + os.Exit(1) + } } func runSetupServer() { @@ -139,13 +235,13 @@ func runSetupServer() { } } -func runMainServer() { +func runMainServer() error { cfg, err := config.LoadForBootstrap() if err != nil { - log.Fatalf("Failed to load config: %v", err) + return fmt.Errorf("load config: %w", err) } if err := logger.Init(logger.OptionsFromConfig(cfg.Log)); err != nil { - log.Fatalf("Failed to initialize logger: %v", err) + return fmt.Errorf("initialize logger: %w", err) } if cfg.RunMode == config.RunModeSimple { log.Println("⚠️ WARNING: Running in SIMPLE mode - billing and quota checks are DISABLED") @@ -155,49 +251,61 @@ func runMainServer() { Version: Version, BuildType: BuildType, } - - app, err := initializeApplication(buildInfo) + listenSpec, err := cfg.Server.ListenSpec() if err != nil { - log.Fatalf("Failed to initialize application: %v", err) + return fmt.Errorf("invalid server listen configuration: %w", err) } - defer app.Cleanup() - pprofServer := startPprofServer() - listenSpec, err := cfg.Server.ListenSpec() + app, err := initializeApplication(buildInfo) if err != nil { - log.Fatalf("Invalid server listen configuration: %v", err) + return fmt.Errorf("initialize application: %w", err) } - // 启动服务器 - go func() { - if err := serveServer(app.Server, listenSpec); err != nil && !errors.Is(err, http.ErrServerClosed) { - log.Fatalf("Failed to start server: %v", err) - } - }() - - log.Printf("Server started on %s", listenSpec.DisplayAddress()) - - // 等待中断信号 - quit := make(chan os.Signal, 1) - signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM) - <-quit + shutdownContext, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM) + defer stop() + restartContext, stopRestart := signal.NotifyContext(context.Background(), syscall.SIGHUP) + defer stopRestart() - log.Println("Shutting down server...") - - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - defer cancel() + serveResults := make(chan serverServeResult, 2) + pprofServer, pprofStartErr := startPprofServer(serveResults) - if err := app.Server.Shutdown(ctx); err != nil { - log.Fatalf("Server forced to shutdown: %v", err) + shutdownTargets := []shutdownTarget{{ + name: "main server", + server: app.Server, + timeout: cfg.Server.ShutdownTimeout(), + }} + if pprofServer != nil { + shutdownTargets = append(shutdownTargets, shutdownTarget{ + name: "pprof server", + server: pprofServer, + timeout: pprofShutdownTimeout, + }) } - if pprofServer != nil { - if err := pprofServer.Shutdown(ctx); err != nil { - log.Fatalf("pprof server forced to shutdown: %v", err) - } + if pprofStartErr != nil { + serveResults <- serverServeResult{name: "pprof server", err: pprofStartErr} + } else { + go func() { + serveResults <- serverServeResult{ + name: "main server", + err: serveServer(app.Server, listenSpec), + } + }() + log.Printf("Server started on %s", listenSpec.DisplayAddress()) } - log.Println("Server exited") + err = runServerLifecycle( + shutdownContext.Done(), + restartContext.Done(), + serveResults, + shutdownTargets, + app.Cleanup, + cfg.Server.ShutdownTimeout(), + ) + if err == nil { + log.Println("Server exited") + } + return err } func serveServer(server *http.Server, spec config.ServerListenSpec) error { @@ -229,18 +337,21 @@ func serveServer(server *http.Server, spec config.ServerListenSpec) error { } } -func startPprofServer() *http.Server { +func startPprofServer(serveResults chan<- serverServeResult) (*http.Server, error) { enabledValue := strings.TrimSpace(os.Getenv("PPROF_ENABLED")) if enabledValue == "" { - return nil + return nil, nil } enabled, err := strconv.ParseBool(enabledValue) if err != nil { - log.Fatalf("Invalid PPROF_ENABLED value %q: %v", enabledValue, err) + return nil, fmt.Errorf("invalid PPROF_ENABLED value %q: %w", enabledValue, err) } if !enabled { - return nil + return nil, nil + } + if serveResults == nil { + return nil, errors.New("start pprof server: nil serve result channel") } addr := strings.TrimSpace(os.Getenv("PPROF_ADDR")) @@ -255,12 +366,17 @@ func startPprofServer() *http.Server { IdleTimeout: 30 * time.Second, } + listener, err := net.Listen("tcp", addr) + if err != nil { + return nil, fmt.Errorf("listen for pprof on %s: %w", addr, err) + } go func() { - if err := server.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) { - log.Fatalf("Failed to start pprof server on %s: %v", addr, err) + serveResults <- serverServeResult{ + name: "pprof server", + err: server.Serve(listener), } }() log.Printf("pprof server started on %s", addr) - return server + return server, nil } diff --git a/backend/cmd/server/main_lifecycle_test.go b/backend/cmd/server/main_lifecycle_test.go new file mode 100644 index 000000000..7fcb573e5 --- /dev/null +++ b/backend/cmd/server/main_lifecycle_test.go @@ -0,0 +1,363 @@ +package main + +import ( + "context" + "errors" + "net/http" + "sync/atomic" + "testing" + "time" +) + +func TestRunServerLifecycleStopRunsShutdownAndCleanup(t *testing.T) { + stop := make(chan struct{}) + close(stop) + restart := make(chan struct{}) + serveResults := make(chan serverServeResult) + + var lifecycleStage atomic.Int32 + server := &shutdownServerStub{ + shutdownFn: func(context.Context) error { + if !lifecycleStage.CompareAndSwap(0, 1) { + lifecycleStage.Store(-1) + } + return nil + }, + } + var cleanupCalls atomic.Int32 + cleanup := func(ctx context.Context) error { + cleanupCalls.Add(1) + if ctx == nil { + lifecycleStage.Store(-1) + return nil + } + if !lifecycleStage.CompareAndSwap(1, 2) { + lifecycleStage.Store(-1) + } + return nil + } + + err := runServerLifecycle( + stop, + restart, + serveResults, + []shutdownTarget{{name: "main", server: server, timeout: time.Second}}, + cleanup, + time.Second, + ) + if err != nil { + t.Fatalf("runServerLifecycle() error = %v", err) + } + if got := server.shutdownCalls.Load(); got != 1 { + t.Fatalf("Shutdown() calls = %d, want 1", got) + } + if got := server.closeCalls.Load(); got != 0 { + t.Fatalf("Close() calls = %d, want 0", got) + } + if got := cleanupCalls.Load(); got != 1 { + t.Fatalf("cleanup calls = %d, want 1", got) + } + if got := lifecycleStage.Load(); got != 2 { + t.Fatalf("lifecycle stage = %d, want 2 (shutdown before cleanup)", got) + } +} + +func TestRunServerLifecycleServeFailureShutsDownAndReturnsCause(t *testing.T) { + serveErr := errors.New("serve failed") + stop := make(chan struct{}) + restart := make(chan struct{}) + serveResults := make(chan serverServeResult, 1) + serveResults <- serverServeResult{name: "main", err: serveErr} + + server := &shutdownServerStub{} + var cleanupCalls atomic.Int32 + err := runServerLifecycle( + stop, + restart, + serveResults, + []shutdownTarget{{name: "main", server: server, timeout: time.Second}}, + func(context.Context) error { + cleanupCalls.Add(1) + return nil + }, + time.Second, + ) + + if !errors.Is(err, serveErr) { + t.Fatalf("runServerLifecycle() error = %v, want wrapping %v", err, serveErr) + } + if got := server.shutdownCalls.Load(); got != 1 { + t.Fatalf("Shutdown() calls = %d, want 1", got) + } + if got := cleanupCalls.Load(); got != 1 { + t.Fatalf("cleanup calls = %d, want 1", got) + } +} + +func TestRunServerLifecycleStopIgnoresShutdownFailureButStillCleansUp(t *testing.T) { + stop := make(chan struct{}) + close(stop) + restart := make(chan struct{}) + serveResults := make(chan serverServeResult) + shutdownErr := errors.New("graceful shutdown failed") + + server := &shutdownServerStub{ + shutdownFn: func(context.Context) error { return shutdownErr }, + } + var cleanupCalls atomic.Int32 + err := runServerLifecycle( + stop, + restart, + serveResults, + []shutdownTarget{{name: "main", server: server, timeout: time.Second}}, + func(context.Context) error { + cleanupCalls.Add(1) + return nil + }, + time.Second, + ) + + if err != nil { + t.Fatalf("runServerLifecycle() error = %v, want nil for requested stop", err) + } + if got := server.shutdownCalls.Load(); got != 1 { + t.Fatalf("Shutdown() calls = %d, want 1", got) + } + if got := server.closeCalls.Load(); got != 1 { + t.Fatalf("Close() calls = %d, want 1 after shutdown failure", got) + } + if got := cleanupCalls.Load(); got != 1 { + t.Fatalf("cleanup calls = %d, want 1", got) + } +} + +func TestRunServerLifecycleStopIgnoresCleanupFailure(t *testing.T) { + stop := make(chan struct{}) + close(stop) + restart := make(chan struct{}) + serveResults := make(chan serverServeResult) + cleanupErr := errors.New("cleanup failed") + + server := &shutdownServerStub{} + var cleanupCalls atomic.Int32 + err := runServerLifecycle( + stop, + restart, + serveResults, + []shutdownTarget{{name: "main", server: server, timeout: time.Second}}, + func(context.Context) error { + cleanupCalls.Add(1) + return cleanupErr + }, + time.Second, + ) + + if err != nil { + t.Fatalf("runServerLifecycle() error = %v, want nil for requested stop", err) + } + if got := server.shutdownCalls.Load(); got != 1 { + t.Fatalf("Shutdown() calls = %d, want 1", got) + } + if got := cleanupCalls.Load(); got != 1 { + t.Fatalf("cleanup calls = %d, want 1", got) + } +} + +func TestRunServerLifecycleRestartRunsShutdownAndCleanup(t *testing.T) { + stop := make(chan struct{}) + restart := make(chan struct{}) + close(restart) + serveResults := make(chan serverServeResult) + + server := &shutdownServerStub{} + var cleanupCalls atomic.Int32 + err := runServerLifecycle( + stop, + restart, + serveResults, + []shutdownTarget{{name: "main", server: server, timeout: time.Second}}, + func(context.Context) error { + cleanupCalls.Add(1) + return nil + }, + time.Second, + ) + + if !errors.Is(err, errServerRestartRequested) { + t.Fatalf("runServerLifecycle() error = %v, want wrapping %v", err, errServerRestartRequested) + } + if got := server.shutdownCalls.Load(); got != 1 { + t.Fatalf("Shutdown() calls = %d, want 1", got) + } + if got := cleanupCalls.Load(); got != 1 { + t.Fatalf("cleanup calls = %d, want 1", got) + } +} + +func TestRunServerLifecycleStopTakesPriorityOverServeFailure(t *testing.T) { + stop := make(chan struct{}) + close(stop) + restart := make(chan struct{}) + serveResults := make(chan serverServeResult, 1) + serveResults <- serverServeResult{name: "main", err: errors.New("serve failed")} + + server := &shutdownServerStub{} + err := runServerLifecycle( + stop, + restart, + serveResults, + []shutdownTarget{{name: "main", server: server, timeout: time.Second}}, + func(context.Context) error { return nil }, + time.Second, + ) + + if err != nil { + t.Fatalf("runServerLifecycle() error = %v, want nil when stop and serve result are both ready", err) + } + if got := server.shutdownCalls.Load(); got != 1 { + t.Fatalf("Shutdown() calls = %d, want 1", got) + } +} + +func TestRunServerLifecycleStopTakesPriorityOverRestart(t *testing.T) { + stop := make(chan struct{}) + close(stop) + restart := make(chan struct{}) + close(restart) + serveResults := make(chan serverServeResult) + + server := &shutdownServerStub{} + err := runServerLifecycle( + stop, + restart, + serveResults, + []shutdownTarget{{name: "main", server: server, timeout: time.Second}}, + func(context.Context) error { return nil }, + time.Second, + ) + + if err != nil { + t.Fatalf("runServerLifecycle() error = %v, want nil when stop and restart are both ready", err) + } + if got := server.shutdownCalls.Load(); got != 1 { + t.Fatalf("Shutdown() calls = %d, want 1", got) + } +} + +func TestRunServerLifecycleRestartTakesPriorityOverServeFailure(t *testing.T) { + stop := make(chan struct{}) + restart := make(chan struct{}) + close(restart) + serveResults := make(chan serverServeResult, 1) + serveResults <- serverServeResult{name: "main", err: errors.New("serve failed")} + + server := &shutdownServerStub{} + var cleanupCalls atomic.Int32 + err := runServerLifecycle( + stop, + restart, + serveResults, + []shutdownTarget{{name: "main", server: server, timeout: time.Second}}, + func(context.Context) error { + cleanupCalls.Add(1) + return nil + }, + time.Second, + ) + + if !errors.Is(err, errServerRestartRequested) { + t.Fatalf("runServerLifecycle() error = %v, want wrapping %v", err, errServerRestartRequested) + } + if got := server.shutdownCalls.Load(); got != 1 { + t.Fatalf("Shutdown() calls = %d, want 1", got) + } + if got := cleanupCalls.Load(); got != 1 { + t.Fatalf("cleanup calls = %d, want 1", got) + } +} + +func TestRunServerLifecycleStopTakesPriorityWhenAllTriggersReady(t *testing.T) { + stop := make(chan struct{}) + close(stop) + restart := make(chan struct{}) + close(restart) + serveResults := make(chan serverServeResult, 1) + serveResults <- serverServeResult{name: "main", err: errors.New("serve failed")} + + server := &shutdownServerStub{} + var cleanupCalls atomic.Int32 + err := runServerLifecycle( + stop, + restart, + serveResults, + []shutdownTarget{{name: "main", server: server, timeout: time.Second}}, + func(context.Context) error { + cleanupCalls.Add(1) + return nil + }, + time.Second, + ) + + if err != nil { + t.Fatalf("runServerLifecycle() error = %v, want nil when all triggers are ready", err) + } + if got := server.shutdownCalls.Load(); got != 1 { + t.Fatalf("Shutdown() calls = %d, want 1", got) + } + if got := cleanupCalls.Load(); got != 1 { + t.Fatalf("cleanup calls = %d, want 1", got) + } +} + +func TestStartPprofServerRejectsInvalidEnabledValue(t *testing.T) { + t.Setenv("PPROF_ENABLED", "not-a-boolean") + serveResults := make(chan serverServeResult, 1) + + server, err := startPprofServer(serveResults) + if err == nil { + t.Fatal("startPprofServer() error = nil, want invalid PPROF_ENABLED error") + } + if server != nil { + t.Fatal("startPprofServer() server is non-nil after invalid PPROF_ENABLED") + } + select { + case result := <-serveResults: + t.Fatalf("startPprofServer() unexpectedly started a server: %+v", result) + default: + } +} + +func TestStartPprofServerReportsGracefulShutdown(t *testing.T) { + t.Setenv("PPROF_ENABLED", "true") + t.Setenv("PPROF_ADDR", "127.0.0.1:0") + serveResults := make(chan serverServeResult, 1) + + server, err := startPprofServer(serveResults) + if err != nil { + t.Fatalf("startPprofServer() error = %v", err) + } + if server == nil { + t.Fatal("startPprofServer() server = nil, want running server") + } + defer func() { + _ = server.Close() + }() + + shutdownCtx, cancelShutdown := context.WithTimeout(context.Background(), time.Second) + defer cancelShutdown() + if err := server.Shutdown(shutdownCtx); err != nil { + t.Fatalf("pprof server Shutdown() error = %v", err) + } + + select { + case result := <-serveResults: + if result.name != "pprof server" { + t.Errorf("serve result name = %q, want %q", result.name, "pprof server") + } + if !errors.Is(result.err, http.ErrServerClosed) { + t.Errorf("serve result error = %v, want wrapping %v", result.err, http.ErrServerClosed) + } + case <-time.After(time.Second): + t.Fatal("timed out waiting for pprof serve result after shutdown") + } +} diff --git a/backend/cmd/server/main_migrate_only_test.go b/backend/cmd/server/main_migrate_only_test.go new file mode 100644 index 000000000..c27e087f6 --- /dev/null +++ b/backend/cmd/server/main_migrate_only_test.go @@ -0,0 +1,170 @@ +package main + +import ( + "context" + "errors" + "strings" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" +) + +func TestCommandOptionsValidateMigrateOnly(t *testing.T) { + tests := []struct { + name string + options commandOptions + wantErr string + }{ + { + name: "valid", + options: commandOptions{ + migrateOnly: true, + migrationTimeout: defaultMigrationTimeout, + }, + }, + { + name: "setup conflict", + options: commandOptions{ + setupMode: true, + migrateOnly: true, + migrationTimeout: defaultMigrationTimeout, + }, + wantErr: "cannot be combined with --setup", + }, + { + name: "version conflict", + options: commandOptions{ + showVersion: true, + migrateOnly: true, + migrationTimeout: defaultMigrationTimeout, + }, + wantErr: "cannot be combined with --version", + }, + { + name: "zero timeout", + options: commandOptions{ + migrateOnly: true, + }, + wantErr: "must be greater than zero", + }, + { + name: "negative timeout", + options: commandOptions{ + migrateOnly: true, + migrationTimeout: -time.Second, + }, + wantErr: "must be greater than zero", + }, + { + name: "timeout ignored without migrate mode", + options: commandOptions{}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := tt.options.validate() + if tt.wantErr == "" { + if err != nil { + t.Fatalf("validate() error = %v", err) + } + return + } + if err == nil || !strings.Contains(err.Error(), tt.wantErr) { + t.Fatalf("validate() error = %v, want containing %q", err, tt.wantErr) + } + }) + } +} + +func TestRunMigrationsOnlyUsesLoadedConfigAndDeadline(t *testing.T) { + wantConfig := &config.Config{} + var gotConfig *config.Config + var gotDeadline bool + + err := runMigrationsOnly( + context.Background(), + time.Minute, + func() (*config.Config, error) { + return wantConfig, nil + }, + func(ctx context.Context, cfg *config.Config) error { + gotConfig = cfg + _, gotDeadline = ctx.Deadline() + return nil + }, + ) + if err != nil { + t.Fatalf("runMigrationsOnly() error = %v", err) + } + if gotConfig != wantConfig { + t.Fatalf("migration config pointer = %p, want %p", gotConfig, wantConfig) + } + if !gotDeadline { + t.Fatal("migration context has no deadline") + } +} + +func TestRunMigrationsOnlyRejectsInvalidDependencies(t *testing.T) { + validLoader := func() (*config.Config, error) { return &config.Config{}, nil } + validRunner := func(context.Context, *config.Config) error { return nil } + var nilContext context.Context + + tests := []struct { + name string + ctx context.Context + timeout time.Duration + loadConfig bootstrapConfigLoader + runner configuredMigrationRunner + wantErr string + }{ + {name: "nil context", ctx: nilContext, timeout: time.Minute, loadConfig: validLoader, runner: validRunner, wantErr: "nil parent context"}, + {name: "zero timeout", ctx: context.Background(), loadConfig: validLoader, runner: validRunner, wantErr: "must be greater than zero"}, + {name: "nil loader", ctx: context.Background(), timeout: time.Minute, runner: validRunner, wantErr: "nil bootstrap config loader"}, + {name: "nil runner", ctx: context.Background(), timeout: time.Minute, loadConfig: validLoader, wantErr: "nil configured migration runner"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := runMigrationsOnly(tt.ctx, tt.timeout, tt.loadConfig, tt.runner) + if err == nil || !strings.Contains(err.Error(), tt.wantErr) { + t.Fatalf("runMigrationsOnly() error = %v, want containing %q", err, tt.wantErr) + } + }) + } +} + +func TestRunMigrationsOnlyStopsOnConfigError(t *testing.T) { + wantErr := errors.New("config unavailable") + runnerCalled := false + + err := runMigrationsOnly( + context.Background(), + time.Minute, + func() (*config.Config, error) { return nil, wantErr }, + func(context.Context, *config.Config) error { + runnerCalled = true + return nil + }, + ) + if !errors.Is(err, wantErr) { + t.Fatalf("runMigrationsOnly() error = %v, want wrapping %v", err, wantErr) + } + if runnerCalled { + t.Fatal("migration runner called after config load failure") + } +} + +func TestRunMigrationsOnlyPropagatesMigrationError(t *testing.T) { + wantErr := context.DeadlineExceeded + err := runMigrationsOnly( + context.Background(), + time.Minute, + func() (*config.Config, error) { return &config.Config{}, nil }, + func(context.Context, *config.Config) error { return wantErr }, + ) + if !errors.Is(err, wantErr) { + t.Fatalf("runMigrationsOnly() error = %v, want wrapping %v", err, wantErr) + } +} diff --git a/backend/cmd/server/main_shutdown_test.go b/backend/cmd/server/main_shutdown_test.go new file mode 100644 index 000000000..9b46d17b7 --- /dev/null +++ b/backend/cmd/server/main_shutdown_test.go @@ -0,0 +1,191 @@ +package main + +import ( + "context" + "errors" + "sync/atomic" + "testing" + "time" +) + +type shutdownServerStub struct { + shutdownFn func(context.Context) error + closeFn func() error + shutdownCalls atomic.Int32 + closeCalls atomic.Int32 +} + +func (s *shutdownServerStub) Shutdown(ctx context.Context) error { + s.shutdownCalls.Add(1) + if s.shutdownFn == nil { + return nil + } + return s.shutdownFn(ctx) +} + +func (s *shutdownServerStub) Close() error { + s.closeCalls.Add(1) + if s.closeFn == nil { + return nil + } + return s.closeFn() +} + +func TestShutdownHTTPServersSuccessfulShutdownDoesNotClose(t *testing.T) { + server := &shutdownServerStub{} + + err := shutdownHTTPServers(shutdownTarget{ + name: "api", + server: server, + timeout: time.Second, + }) + if err != nil { + t.Fatalf("shutdownHTTPServers() error = %v", err) + } + if got := server.shutdownCalls.Load(); got != 1 { + t.Fatalf("Shutdown() calls = %d, want 1", got) + } + if got := server.closeCalls.Load(); got != 0 { + t.Fatalf("Close() calls = %d, want 0", got) + } +} + +func TestShutdownHTTPServersFailureForcesClose(t *testing.T) { + wantErr := errors.New("graceful shutdown failed") + server := &shutdownServerStub{ + shutdownFn: func(context.Context) error { return wantErr }, + } + + err := shutdownHTTPServers(shutdownTarget{ + name: "api", + server: server, + timeout: time.Second, + }) + if !errors.Is(err, wantErr) { + t.Fatalf("shutdownHTTPServers() error = %v, want wrapping %v", err, wantErr) + } + if got := server.closeCalls.Load(); got != 1 { + t.Fatalf("Close() calls = %d, want 1", got) + } +} + +func TestShutdownHTTPServersDeadlineForcesClose(t *testing.T) { + server := &shutdownServerStub{ + shutdownFn: func(ctx context.Context) error { + <-ctx.Done() + return ctx.Err() + }, + } + + err := shutdownHTTPServers(shutdownTarget{ + name: "api", + server: server, + timeout: 20 * time.Millisecond, + }) + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("shutdownHTTPServers() error = %v, want wrapping %v", err, context.DeadlineExceeded) + } + if got := server.closeCalls.Load(); got != 1 { + t.Fatalf("Close() calls = %d, want 1", got) + } +} + +func TestShutdownHTTPServersPreservesShutdownAndCloseErrors(t *testing.T) { + shutdownErr := errors.New("shutdown failed") + closeErr := errors.New("forced close failed") + server := &shutdownServerStub{ + shutdownFn: func(context.Context) error { return shutdownErr }, + closeFn: func() error { return closeErr }, + } + + err := shutdownHTTPServers(shutdownTarget{ + name: "api", + server: server, + timeout: time.Second, + }) + if !errors.Is(err, shutdownErr) { + t.Errorf("shutdownHTTPServers() error = %v, want wrapping shutdown error %v", err, shutdownErr) + } + if !errors.Is(err, closeErr) { + t.Errorf("shutdownHTTPServers() error = %v, want wrapping close error %v", err, closeErr) + } +} + +func TestShutdownHTTPServersContinuesAfterTargetFailure(t *testing.T) { + firstErr := errors.New("first server failed") + first := &shutdownServerStub{ + shutdownFn: func(context.Context) error { return firstErr }, + } + second := &shutdownServerStub{} + + err := shutdownHTTPServers( + shutdownTarget{name: "api", server: first, timeout: time.Second}, + shutdownTarget{name: "pprof", server: second, timeout: time.Second}, + ) + if !errors.Is(err, firstErr) { + t.Fatalf("shutdownHTTPServers() error = %v, want wrapping %v", err, firstErr) + } + if got := first.closeCalls.Load(); got != 1 { + t.Fatalf("first Close() calls = %d, want 1", got) + } + if got := second.shutdownCalls.Load(); got != 1 { + t.Fatalf("second Shutdown() calls = %d, want 1", got) + } + if got := second.closeCalls.Load(); got != 0 { + t.Fatalf("second Close() calls = %d, want 0", got) + } +} + +func TestShutdownHTTPServersUsesIndependentContexts(t *testing.T) { + first := &shutdownServerStub{ + shutdownFn: func(ctx context.Context) error { + <-ctx.Done() + return ctx.Err() + }, + } + secondContextWasActive := atomic.Bool{} + second := &shutdownServerStub{ + shutdownFn: func(ctx context.Context) error { + select { + case <-ctx.Done(): + return errors.New("second server received an expired context") + default: + secondContextWasActive.Store(true) + return nil + } + }, + } + + err := shutdownHTTPServers( + shutdownTarget{name: "api", server: first, timeout: 20 * time.Millisecond}, + shutdownTarget{name: "pprof", server: second, timeout: time.Second}, + ) + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("shutdownHTTPServers() error = %v, want wrapping %v", err, context.DeadlineExceeded) + } + if !secondContextWasActive.Load() { + t.Fatal("second server did not receive a fresh active context") + } + if got := second.closeCalls.Load(); got != 0 { + t.Fatalf("second Close() calls = %d, want 0", got) + } +} + +func TestShutdownHTTPServersRejectsInvalidTimeoutWithoutCallingServer(t *testing.T) { + server := &shutdownServerStub{} + + err := shutdownHTTPServers(shutdownTarget{ + name: "api", + server: server, + timeout: 0, + }) + if err == nil { + t.Fatal("shutdownHTTPServers() error = nil, want invalid timeout error") + } + if got := server.shutdownCalls.Load(); got != 0 { + t.Fatalf("Shutdown() calls = %d, want 0", got) + } + if got := server.closeCalls.Load(); got != 0 { + t.Fatalf("Close() calls = %d, want 0", got) + } +} diff --git a/backend/cmd/server/shutdown.go b/backend/cmd/server/shutdown.go new file mode 100644 index 000000000..fd4d4dcf3 --- /dev/null +++ b/backend/cmd/server/shutdown.go @@ -0,0 +1,288 @@ +package main + +import ( + "context" + "errors" + "fmt" + "log" + "sort" + "strings" + "time" +) + +const pprofShutdownTimeout = 5 * time.Second + +var errServerRestartRequested = errors.New("service restart requested") + +type shutdownHTTPServer interface { + Shutdown(context.Context) error + Close() error +} + +type shutdownTarget struct { + name string + server shutdownHTTPServer + timeout time.Duration +} + +type serverServeResult struct { + name string + err error +} + +// runServerLifecycle waits for an operator stop or an unexpected server exit, +// then always runs HTTP shutdown and application cleanup. Errors encountered +// during an operator-initiated stop are logged but do not turn a normal +// systemd stop into a failed process exit. +func runServerLifecycle( + stop <-chan struct{}, + restart <-chan struct{}, + serveResults <-chan serverServeResult, + targets []shutdownTarget, + cleanup func(context.Context) error, + cleanupTimeout time.Duration, +) error { + if stop == nil { + return errors.New("server lifecycle: nil stop channel") + } + if restart == nil { + return errors.New("server lifecycle: nil restart channel") + } + if serveResults == nil { + return errors.New("server lifecycle: nil serve result channel") + } + if cleanup == nil { + return errors.New("server lifecycle: nil cleanup function") + } + if cleanupTimeout <= 0 { + return errors.New("server lifecycle: cleanup timeout must be greater than zero") + } + + lifecycleErr := waitForServerLifecycleTrigger(stop, restart, serveResults) + switch { + case lifecycleErr == nil: + log.Println("Shutting down server...") + case errors.Is(lifecycleErr, errServerRestartRequested): + log.Println("Graceful service restart requested...") + default: + log.Printf("Server lifecycle error: %v", lifecycleErr) + } + + if err := shutdownHTTPServers(targets...); err != nil { + log.Printf("Server shutdown completed with forced-close errors: %v", err) + } + + cleanupCtx, cancelCleanup := context.WithTimeout(context.Background(), cleanupTimeout) + cleanupErr := cleanup(cleanupCtx) + cancelCleanup() + if cleanupErr != nil { + log.Printf("Application cleanup completed with errors: %v", cleanupErr) + } + + return lifecycleErr +} + +func waitForServerLifecycleTrigger( + stop <-chan struct{}, + restart <-chan struct{}, + serveResults <-chan serverServeResult, +) error { + select { + case <-stop: + return nil + case <-restart: + if channelReady(stop) { + return nil + } + return errServerRestartRequested + case result := <-serveResults: + // Operator stop has the highest priority, followed by an explicit + // restart request. This removes random exit-code selection when more + // than one channel becomes ready at the same time. + if channelReady(stop) { + return nil + } + if channelReady(restart) { + return errServerRestartRequested + } + + name := strings.TrimSpace(result.name) + if name == "" { + name = "HTTP server" + } + if result.err == nil { + return fmt.Errorf("%s stopped unexpectedly without an error", name) + } + return fmt.Errorf("%s stopped unexpectedly: %w", name, result.err) + } +} + +func channelReady(ch <-chan struct{}) bool { + select { + case <-ch: + return true + default: + return false + } +} + +// shutdownHTTPServers gives each server an independent graceful-shutdown +// budget. If graceful shutdown fails, Close is used to release HTTP/SSE +// connections before the next target is processed. +func shutdownHTTPServers(targets ...shutdownTarget) error { + var shutdownErrors []error + for _, target := range targets { + name := strings.TrimSpace(target.name) + if name == "" { + name = "HTTP server" + } + if target.server == nil { + shutdownErrors = append(shutdownErrors, fmt.Errorf("%s shutdown: nil server", name)) + continue + } + if target.timeout <= 0 { + shutdownErrors = append(shutdownErrors, fmt.Errorf("%s shutdown: timeout must be greater than zero", name)) + continue + } + + ctx, cancel := context.WithTimeout(context.Background(), target.timeout) + shutdownErr := target.server.Shutdown(ctx) + cancel() + if shutdownErr == nil { + continue + } + + shutdownErr = fmt.Errorf("%s graceful shutdown: %w", name, shutdownErr) + if closeErr := target.server.Close(); closeErr != nil { + shutdownErrors = append(shutdownErrors, errors.Join( + shutdownErr, + fmt.Errorf("%s forced close: %w", name, closeErr), + )) + continue + } + shutdownErrors = append(shutdownErrors, shutdownErr) + } + return errors.Join(shutdownErrors...) +} + +type cleanupStep struct { + name string + fn func() error +} + +type cleanupStepResult struct { + index int + err error +} + +// runProcessCleanup enforces a process-exit deadline around cleanup +// orchestration. A blocking step is not context-cancelled; on timeout the +// caller must return from the main process so the runtime can reclaim it. +func runProcessCleanup(ctx context.Context, parallelSteps, infraSteps []cleanupStep) error { + if ctx == nil { + return errors.New("process cleanup: nil context") + } + if err := validateCleanupSteps(parallelSteps, infraSteps); err != nil { + return err + } + + var cleanupErrors []error + results := make(chan cleanupStepResult, len(parallelSteps)) + pending := make(map[int]string, len(parallelSteps)) + for index, step := range parallelSteps { + pending[index] = step.name + go func(index int, step cleanupStep) { + results <- cleanupStepResult{index: index, err: step.fn()} + }(index, step) + } + + for len(pending) > 0 { + select { + case result := <-results: + name := pending[result.index] + delete(pending, result.index) + if result.err != nil { + wrapped := fmt.Errorf("cleanup %s: %w", name, result.err) + cleanupErrors = append(cleanupErrors, wrapped) + log.Printf("[Cleanup] %s failed: %v", name, result.err) + continue + } + log.Printf("[Cleanup] %s succeeded", name) + case <-ctx.Done(): + return errors.Join( + append(cleanupErrors, cleanupDeadlineError(ctx.Err(), pendingCleanupNames(pending)))..., + ) + } + } + + for index, step := range infraSteps { + if err := ctx.Err(); err != nil { + return errors.Join( + append(cleanupErrors, cleanupDeadlineError(err, cleanupStepNames(infraSteps[index:])))..., + ) + } + + result := make(chan error, 1) + go func(step cleanupStep) { + result <- step.fn() + }(step) + select { + case err := <-result: + if err != nil { + wrapped := fmt.Errorf("cleanup %s: %w", step.name, err) + cleanupErrors = append(cleanupErrors, wrapped) + log.Printf("[Cleanup] %s failed: %v", step.name, err) + continue + } + log.Printf("[Cleanup] %s succeeded", step.name) + case <-ctx.Done(): + return errors.Join( + append(cleanupErrors, cleanupDeadlineError(ctx.Err(), cleanupStepNames(infraSteps[index:])))..., + ) + } + } + + if len(cleanupErrors) == 0 { + log.Printf("[Cleanup] All cleanup steps completed") + } + return errors.Join(cleanupErrors...) +} + +func validateCleanupSteps(stepGroups ...[]cleanupStep) error { + for _, steps := range stepGroups { + for index, step := range steps { + if strings.TrimSpace(step.name) == "" { + return fmt.Errorf("process cleanup: step %d has an empty name", index) + } + if step.fn == nil { + return fmt.Errorf("process cleanup: step %s has a nil function", step.name) + } + } + } + return nil +} + +func cleanupDeadlineError(cause error, pending []string) error { + return fmt.Errorf( + "process cleanup deadline reached with pending steps [%s]: %w", + strings.Join(pending, ", "), + cause, + ) +} + +func pendingCleanupNames(pending map[int]string) []string { + names := make([]string, 0, len(pending)) + for _, name := range pending { + names = append(names, name) + } + sort.Strings(names) + return names +} + +func cleanupStepNames(steps []cleanupStep) []string { + names := make([]string, 0, len(steps)) + for _, step := range steps { + names = append(names, step.name) + } + return names +} diff --git a/backend/cmd/server/wire.go b/backend/cmd/server/wire.go index 5abbb2fdf..852c99f74 100644 --- a/backend/cmd/server/wire.go +++ b/backend/cmd/server/wire.go @@ -5,10 +5,7 @@ package main import ( "context" - "log" "net/http" - "sync" - "time" "github.com/Wei-Shaw/sub2api/ent" "github.com/Wei-Shaw/sub2api/internal/config" @@ -25,7 +22,7 @@ import ( type Application struct { Server *http.Server - Cleanup func() + Cleanup func(context.Context) error } func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { @@ -103,16 +100,8 @@ func provideCleanup( activityAutoDraw *service.ActivityAutoDrawService, paymentOrderExpiry *service.PaymentOrderExpiryService, channelMonitorRunner *service.ChannelMonitorRunner, -) func() { - return func() { - ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) - defer cancel() - - type cleanupStep struct { - name string - fn func() error - } - +) func(context.Context) error { + return func(ctx context.Context) error { // 应用层清理步骤可并行执行，基础设施资源（Redis/Ent）最后按顺序关闭。 parallelSteps := []cleanupStep{ {"OpsScheduledReportService", func() error { @@ -299,43 +288,6 @@ func provideCleanup( }}, } - runParallel := func(steps []cleanupStep) { - var wg sync.WaitGroup - for i := range steps { - step := steps[i] - wg.Add(1) - go func() { - defer wg.Done() - if err := step.fn(); err != nil { - log.Printf("[Cleanup] %s failed: %v", step.name, err) - return - } - log.Printf("[Cleanup] %s succeeded", step.name) - }() - } - wg.Wait() - } - - runSequential := func(steps []cleanupStep) { - for i := range steps { - step := steps[i] - if err := step.fn(); err != nil { - log.Printf("[Cleanup] %s failed: %v", step.name, err) - continue - } - log.Printf("[Cleanup] %s succeeded", step.name) - } - } - - runParallel(parallelSteps) - runSequential(infraSteps) - - // Check if context timed out - select { - case <-ctx.Done(): - log.Printf("[Cleanup] Warning: cleanup timed out after 10 seconds") - default: - log.Printf("[Cleanup] All cleanup steps completed") - } + return runProcessCleanup(ctx, parallelSteps, infraSteps) } } diff --git a/backend/cmd/server/wire_gen.go b/backend/cmd/server/wire_gen.go index ba75e5d86..24a5708cb 100644 --- a/backend/cmd/server/wire_gen.go +++ b/backend/cmd/server/wire_gen.go @@ -18,10 +18,7 @@ import ( "github.com/Wei-Shaw/sub2api/internal/server/middleware" "github.com/Wei-Shaw/sub2api/internal/service" "github.com/redis/go-redis/v9" - "log" "net/http" - "sync" - "time" ) import ( @@ -147,7 +144,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { accountShareModeService := service.ProvideAccountShareModeService(configConfig, accountShareModeRepository, accountRepository, apiKeyRepository, usageLogRepository, userRepository, proxyRepository, openAIOAuthService, oAuthService, concurrencyService, apiKeyAuthCacheInvalidator, accountTestService, rateLimitService, billingCacheService, billingService, modelPricingResolver, settingRepository, settingService) accountShareModeHandler := handler.NewAccountShareModeHandler(accountShareModeService) accountSharePolicyRepository := repository.NewAccountSharePolicyRepository(client, db) - accountService := service.ProvideAccountService(accountRepository, groupRepository, userRepository, userSubscriptionRepository, proxyRepository, accountSharePolicyRepository, accountShareModeRepository, userPrivateGroupProvisioner, systemNoticeService, settingService) + accountService := service.ProvideAccountService(accountRepository, groupRepository, userRepository, userSubscriptionRepository, proxyRepository, accountSharePolicyRepository, accountShareModeRepository, userPrivateGroupProvisioner, systemNoticeService, settingService, agentIdentityWSInvalidatorProxy) claudeUsageFetcher := repository.NewClaudeUsageFetcher(httpUpstream) antigravityQuotaFetcher := service.NewAntigravityQuotaFetcher(proxyRepository) grokQuotaFetcher := service.NewGrokQuotaFetcher() @@ -192,7 +189,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { dashboardHandler := admin.NewDashboardHandler(dashboardService, dashboardAggregationService) proxyExitInfoProber := repository.NewProxyExitInfoProber(configConfig) proxyLatencyCache := repository.NewProxyLatencyCache(redisClient) - adminService := service.ProvideAdminService(userRepository, groupRepository, accountRepository, proxyRepository, apiKeyRepository, accountShareAPIKeyBindingChecker, redeemCodeRepository, userGroupRateRepository, userRPMCache, billingCacheService, proxyExitInfoProber, proxyLatencyCache, apiKeyAuthCacheInvalidator, client, settingService, subscriptionService, userSubscriptionRepository, privacyClientFactory, userPrivateGroupProvisioner, systemNoticeService) + adminService := service.ProvideAdminService(userRepository, groupRepository, accountRepository, proxyRepository, apiKeyRepository, accountShareAPIKeyBindingChecker, redeemCodeRepository, userGroupRateRepository, userRPMCache, billingCacheService, proxyExitInfoProber, proxyLatencyCache, apiKeyAuthCacheInvalidator, client, settingService, subscriptionService, userSubscriptionRepository, privacyClientFactory, userPrivateGroupProvisioner, systemNoticeService, agentIdentityWSInvalidatorProxy) adminUserHandler := admin.NewUserHandler(adminService, concurrencyService) groupRateScheduleRepository := repository.NewGroupRateScheduleRepository(db) groupRateScheduleService := service.ProvideGroupRateScheduleService(groupRateScheduleRepository, groupRepository, apiKeyAuthCacheInvalidator, apiKeyRepository, userSubscriptionRepository, userGroupRateRepository, systemNoticeService) @@ -294,7 +291,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { userMsgQueueCache := repository.NewUserMsgQueueCache(redisClient) userMessageQueueService := service.ProvideUserMessageQueueService(userMsgQueueCache, rpmCache, configConfig) gatewayHandler := handler.NewGatewayHandler(gatewayService, geminiMessagesCompatService, antigravityGatewayService, userService, concurrencyService, billingCacheService, usageService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, userContentModerationService, userMessageQueueService, configConfig, settingService) - openAIGatewayHandler := handler.NewOpenAIGatewayHandler(openAIGatewayService, concurrencyService, billingCacheService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, userContentModerationService, configConfig) + openAIGatewayHandler := handler.ProvideOpenAIGatewayHandler(openAIGatewayService, concurrencyService, billingCacheService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, userContentModerationService, grokQuotaService, configConfig) handlerSettingHandler := handler.ProvideSettingHandler(settingService, buildInfo) totpHandler := handler.NewTotpHandler(totpService) handlerPaymentHandler := handler.NewPaymentHandler(paymentService, paymentConfigService) @@ -316,7 +313,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { opsMetricsCollector := service.ProvideOpsMetricsCollector(opsRepository, settingRepository, accountRepository, concurrencyService, db, redisClient, configConfig) opsAggregationService := service.ProvideOpsAggregationService(opsRepository, settingRepository, db, redisClient, configConfig) opsAlertEvaluatorService := service.ProvideOpsAlertEvaluatorService(opsService, opsRepository, emailService, redisClient, configConfig) - opsCleanupService := service.ProvideOpsCleanupService(opsRepository, db, redisClient, configConfig, channelMonitorService, backupService) + opsCleanupService := service.ProvideOpsCleanupService(opsService, opsRepository, settingRepository, db, redisClient, configConfig, channelMonitorService, backupService) opsScheduledReportService := service.ProvideOpsScheduledReportService(opsService, userService, emailService, redisClient, configConfig) affiliateCodeCycleService := service.ProvideAffiliateCodeCycleService(affiliateService) tokenRefreshService := service.ProvideTokenRefreshService(accountRepository, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, grokOAuthService, compositeTokenCacheInvalidator, schedulerCache, configConfig, tempUnschedCache, privacyClientFactory, proxyRepository, oAuthRefreshAPI) @@ -338,7 +335,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { type Application struct { Server *http.Server - Cleanup func() + Cleanup func(context.Context) error } func providePrivacyClientFactory() service.PrivacyClientFactory { @@ -386,15 +383,8 @@ func provideCleanup( activityAutoDraw *service.ActivityAutoDrawService, paymentOrderExpiry *service.PaymentOrderExpiryService, channelMonitorRunner *service.ChannelMonitorRunner, -) func() { - return func() { - ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) - defer cancel() - - type cleanupStep struct { - name string - fn func() error - } +) func(context.Context) error { + return func(ctx context.Context) error { parallelSteps := []cleanupStep{ {"OpsScheduledReportService", func() error { @@ -581,42 +571,6 @@ func provideCleanup( }}, } - runParallel := func(steps []cleanupStep) { - var wg sync.WaitGroup - for i := range steps { - step := steps[i] - wg.Add(1) - go func() { - defer wg.Done() - if err := step.fn(); err != nil { - log.Printf("[Cleanup] %s failed: %v", step.name, err) - return - } - log.Printf("[Cleanup] %s succeeded", step.name) - }() - } - wg.Wait() - } - - runSequential := func(steps []cleanupStep) { - for i := range steps { - step := steps[i] - if err := step.fn(); err != nil { - log.Printf("[Cleanup] %s failed: %v", step.name, err) - continue - } - log.Printf("[Cleanup] %s succeeded", step.name) - } - } - - runParallel(parallelSteps) - runSequential(infraSteps) - - select { - case <-ctx.Done(): - log.Printf("[Cleanup] Warning: cleanup timed out after 10 seconds") - default: - log.Printf("[Cleanup] All cleanup steps completed") - } + return runProcessCleanup(ctx, parallelSteps, infraSteps) } } diff --git a/backend/cmd/server/wire_gen_test.go b/backend/cmd/server/wire_gen_test.go index 056006b72..50b5ecd5e 100644 --- a/backend/cmd/server/wire_gen_test.go +++ b/backend/cmd/server/wire_gen_test.go @@ -1,6 +1,7 @@ package main import ( + "context" "testing" "time" @@ -84,7 +85,7 @@ func TestProvideCleanup_WithMinimalDependencies_NoPanic(t *testing.T) { nil, // channelMonitorRunner ) - require.NotPanics(t, func() { - cleanup() - }) + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + require.NoError(t, cleanup(ctx)) } diff --git a/backend/internal/config/config.go b/backend/internal/config/config.go index 17f7e5d8e..195062732 100644 --- a/backend/internal/config/config.go +++ b/backend/internal/config/config.go @@ -8,6 +8,8 @@ import ( "errors" "fmt" "log/slog" + "net" + "net/textproto" "net/url" "os" pathpkg "path" @@ -16,6 +18,7 @@ import ( "time" "github.com/spf13/viper" + "golang.org/x/net/http/httpguts" ) const ( @@ -546,16 +549,31 @@ type PricingConfig struct { } type ServerConfig struct { - Host string `mapstructure:"host"` - Port int `mapstructure:"port"` - Mode string `mapstructure:"mode"` // debug/release - EnableServerTiming bool `mapstructure:"enable_server_timing"` // 已认证管理端/用户端 Web API 性能指标 - FrontendURL string `mapstructure:"frontend_url"` // 前端基础 URL，用于生成邮件中的外部链接 - ReadHeaderTimeout int `mapstructure:"read_header_timeout"` // 读取请求头超时（秒） - IdleTimeout int `mapstructure:"idle_timeout"` // 空闲连接超时（秒） - TrustedProxies []string `mapstructure:"trusted_proxies"` // 可信代理列表（CIDR/IP） - MaxRequestBodySize int64 `mapstructure:"max_request_body_size"` // 全局最大请求体限制 - H2C H2CConfig `mapstructure:"h2c"` // HTTP/2 Cleartext 配置 + Host string `mapstructure:"host"` + Port int `mapstructure:"port"` + Mode string `mapstructure:"mode"` // debug/release + EnableServerTiming bool `mapstructure:"enable_server_timing"` // 已认证管理端/用户端 Web API 性能指标 + FrontendURL string `mapstructure:"frontend_url"` // 前端基础 URL，用于生成邮件中的外部链接 + ReadHeaderTimeout int `mapstructure:"read_header_timeout"` // 读取请求头超时（秒） + IdleTimeout int `mapstructure:"idle_timeout"` // 空闲连接超时（秒） + ShutdownTimeoutSeconds int `mapstructure:"shutdown_timeout_seconds"` // HTTP 服务与应用清理的单阶段退出预算（秒） + TrustedProxies []string `mapstructure:"trusted_proxies"` // 可信代理列表（CIDR/IP） + MaxRequestBodySize int64 `mapstructure:"max_request_body_size"` // 全局最大请求体限制 + H2C H2CConfig `mapstructure:"h2c"` // HTTP/2 Cleartext 配置 +} + +const ( + defaultServerShutdownTimeoutSeconds = 30 + maxServerShutdownTimeoutSeconds = 60 * 60 +) + +// ShutdownTimeout returns the budget used independently by each shutdown phase. +// A zero value is safe for callers that construct ServerConfig directly. +func (c ServerConfig) ShutdownTimeout() time.Duration { + if c.ShutdownTimeoutSeconds == 0 { + return defaultServerShutdownTimeoutSeconds * time.Second + } + return time.Duration(c.ShutdownTimeoutSeconds) * time.Second } type ServerListenNetwork string @@ -707,11 +725,57 @@ type CORSConfig struct { } type SecurityConfig struct { - URLAllowlist URLAllowlistConfig `mapstructure:"url_allowlist"` - ResponseHeaders ResponseHeaderConfig `mapstructure:"response_headers"` - CSP CSPConfig `mapstructure:"csp"` - ProxyFallback ProxyFallbackConfig `mapstructure:"proxy_fallback"` - ProxyProbe ProxyProbeConfig `mapstructure:"proxy_probe"` + URLAllowlist URLAllowlistConfig `mapstructure:"url_allowlist"` + ResponseHeaders ResponseHeaderConfig `mapstructure:"response_headers"` + CSP CSPConfig `mapstructure:"csp"` + ProxyFallback ProxyFallbackConfig `mapstructure:"proxy_fallback"` + ProxyProbe ProxyProbeConfig `mapstructure:"proxy_probe"` + ForwardedClientIPHeaders []string `mapstructure:"forwarded_client_ip_headers"` +} + +const MaxForwardedClientIPHeaders = 16 + +var forbiddenForwardedClientIPHeaders = map[string]struct{}{ + "authorization": {}, + "connection": {}, + "content-length": {}, + "cookie": {}, + "host": {}, + "proxy-authenticate": {}, + "proxy-authorization": {}, + "set-cookie": {}, + "te": {}, + "trailer": {}, + "transfer-encoding": {}, + "upgrade": {}, +} + +// NormalizeForwardedClientIPHeaders 校验并规范化自定义客户端 IP 请求头。 +// 这些请求头仅会在直连来源命中 server.trusted_proxies 时由 Gin 解析。 +func NormalizeForwardedClientIPHeaders(headers []string) ([]string, error) { + normalized := make([]string, 0, len(headers)) + seen := make(map[string]struct{}, len(headers)) + for _, header := range headers { + header = strings.TrimSpace(header) + if !httpguts.ValidHeaderFieldName(header) { + return nil, fmt.Errorf("invalid HTTP header field name %q", header) + } + + canonical := textproto.CanonicalMIMEHeaderKey(header) + key := strings.ToLower(canonical) + if _, forbidden := forbiddenForwardedClientIPHeaders[key]; forbidden { + return nil, fmt.Errorf("HTTP header %q cannot be used as a client IP source", canonical) + } + if _, exists := seen[key]; exists { + continue + } + if len(normalized) >= MaxForwardedClientIPHeaders { + return nil, fmt.Errorf("forwarded client IP headers must contain at most %d unique names", MaxForwardedClientIPHeaders) + } + seen[key] = struct{}{} + normalized = append(normalized, canonical) + } + return normalized, nil } type URLAllowlistConfig struct { @@ -777,10 +841,11 @@ type GatewayConfig struct { // ImageNonstreamTotalTimeoutSeconds: Images 非流式请求的总超时时间（秒），0表示禁用。 // 图片生成可能长时间没有响应体数据，不能复用普通流数据间隔超时。 ImageNonstreamTotalTimeoutSeconds int `mapstructure:"image_nonstream_total_timeout_seconds"` - // OpenAIFirstOutputTimeoutSeconds: OpenAI 原生 HTTP Responses 首个语义输出超时（秒），0 表示禁用。 + // OpenAIFirstOutputTimeoutSeconds: OpenAI HTTP Responses（含 passthrough）首个语义输出超时（秒），0 表示禁用。 + // 默认 60 秒，仅用于截断异常长尾，不限制已开始输出的正常长流。 OpenAIFirstOutputTimeoutSeconds int `mapstructure:"openai_first_output_timeout_seconds"` // OpenAIHighEffortFirstOutputTimeoutSeconds: high/xhigh/max 推理的首个语义输出超时（秒）。 - // 0 表示回退到 OpenAIFirstOutputTimeoutSeconds。 + // 默认 180 秒；0 表示回退到 OpenAIFirstOutputTimeoutSeconds。 OpenAIHighEffortFirstOutputTimeoutSeconds int `mapstructure:"openai_high_effort_first_output_timeout_seconds"` // 请求体最大字节数，用于网关请求体大小限制 MaxBodySize int64 `mapstructure:"max_body_size"` @@ -1349,8 +1414,20 @@ type OpsCleanupConfig struct { Schedule string `mapstructure:"schedule"` // ArchiveExpireDays controls backup record/object expiry for ops log archives. 0 means never expire. ArchiveExpireDays int `mapstructure:"archive_expire_days"` - - // Retention days (0 disables that cleanup target). + // ArchiveWindowDays bounds each ops log archive/delete window to avoid repeatedly exporting the full backlog. + ArchiveWindowDays int `mapstructure:"archive_window_days"` + // MaxCatchupWindowsPerRun limits how many historical windows each ops log table may advance per scheduled run. + MaxCatchupWindowsPerRun int `mapstructure:"max_catchup_windows_per_run"` + // ArchiveTimeoutSeconds bounds one table/window export, compression, and object-storage upload. + ArchiveTimeoutSeconds int `mapstructure:"archive_timeout_seconds"` + // DeleteTimeoutSeconds bounds one table/window batched deletion phase. + DeleteTimeoutSeconds int `mapstructure:"delete_timeout_seconds"` + // RunTimeoutSeconds bounds the complete scheduled cleanup run. + RunTimeoutSeconds int `mapstructure:"run_timeout_seconds"` + // DeleteBatchSize limits rows deleted by each SQL statement. + DeleteBatchSize int `mapstructure:"delete_batch_size"` + + // Retention days. 0 disables that cleanup target. // // vNext requirement: default 30 days across ops datasets. ErrorLogRetentionDays int `mapstructure:"error_log_retention_days"` @@ -1358,6 +1435,50 @@ type OpsCleanupConfig struct { HourlyMetricsRetentionDays int `mapstructure:"hourly_metrics_retention_days"` } +func (c OpsCleanupConfig) Validate() error { + if c.ArchiveExpireDays < 0 { + return fmt.Errorf("ops.cleanup.archive_expire_days must be non-negative") + } + if c.ErrorLogRetentionDays < 0 { + return fmt.Errorf("ops.cleanup.error_log_retention_days must be non-negative") + } + if c.MinuteMetricsRetentionDays < 0 { + return fmt.Errorf("ops.cleanup.minute_metrics_retention_days must be non-negative") + } + if c.HourlyMetricsRetentionDays < 0 { + return fmt.Errorf("ops.cleanup.hourly_metrics_retention_days must be non-negative") + } + if !c.Enabled { + return nil + } + if strings.TrimSpace(c.Schedule) == "" { + return fmt.Errorf("ops.cleanup.schedule is required when ops.cleanup.enabled=true") + } + if c.ArchiveWindowDays <= 0 { + return fmt.Errorf("ops.cleanup.archive_window_days must be positive") + } + if c.MaxCatchupWindowsPerRun <= 0 { + return fmt.Errorf("ops.cleanup.max_catchup_windows_per_run must be positive") + } + if c.ArchiveTimeoutSeconds <= 0 { + return fmt.Errorf("ops.cleanup.archive_timeout_seconds must be positive") + } + if c.DeleteTimeoutSeconds <= 0 { + return fmt.Errorf("ops.cleanup.delete_timeout_seconds must be positive") + } + if c.RunTimeoutSeconds <= 0 { + return fmt.Errorf("ops.cleanup.run_timeout_seconds must be positive") + } + if c.DeleteBatchSize <= 0 { + return fmt.Errorf("ops.cleanup.delete_batch_size must be positive") + } + minimumRunSeconds := int64(c.MaxCatchupWindowsPerRun) * 2 * int64(c.ArchiveTimeoutSeconds+c.DeleteTimeoutSeconds) + if int64(c.RunTimeoutSeconds) < minimumRunSeconds { + return fmt.Errorf("ops.cleanup.run_timeout_seconds must be at least %d for configured catch-up windows", minimumRunSeconds) + } + return nil +} + type OpsAggregationConfig struct { Enabled bool `mapstructure:"enabled"` } @@ -1618,6 +1739,14 @@ func load(allowMissingJWTSecret bool) (*Config, error) { cfg.Security.ResponseHeaders.AdditionalAllowed = normalizeStringSlice(cfg.Security.ResponseHeaders.AdditionalAllowed) cfg.Security.ResponseHeaders.ForceRemove = normalizeStringSlice(cfg.Security.ResponseHeaders.ForceRemove) cfg.Security.CSP.Policy = strings.TrimSpace(cfg.Security.CSP.Policy) + if rawHeaders, configured := os.LookupEnv("SECURITY_FORWARDED_CLIENT_IP_HEADERS"); configured { + cfg.Security.ForwardedClientIPHeaders = normalizeStringSlice(strings.Split(rawHeaders, ",")) + } + forwardedClientIPHeaders, err := NormalizeForwardedClientIPHeaders(cfg.Security.ForwardedClientIPHeaders) + if err != nil { + return nil, fmt.Errorf("security.forwarded_client_ip_headers: %w", err) + } + cfg.Security.ForwardedClientIPHeaders = forwardedClientIPHeaders cfg.Log.Level = strings.ToLower(strings.TrimSpace(cfg.Log.Level)) cfg.Log.Format = strings.ToLower(strings.TrimSpace(cfg.Log.Format)) cfg.Log.ServiceName = strings.TrimSpace(cfg.Log.ServiceName) @@ -1709,6 +1838,7 @@ func setDefaults() { viper.SetDefault("server.frontend_url", "") viper.SetDefault("server.read_header_timeout", 30) // 30秒读取请求头 viper.SetDefault("server.idle_timeout", 120) // 120秒空闲超时 + viper.SetDefault("server.shutdown_timeout_seconds", defaultServerShutdownTimeoutSeconds) viper.SetDefault("server.trusted_proxies", []string{}) viper.SetDefault("server.max_request_body_size", int64(256*1024*1024)) // H2C 默认配置 @@ -1766,6 +1896,7 @@ func setDefaults() { viper.SetDefault("security.csp.enabled", true) viper.SetDefault("security.csp.policy", DefaultCSPPolicy) viper.SetDefault("security.proxy_probe.insecure_skip_verify", false) + viper.SetDefault("security.forwarded_client_ip_headers", []string{}) // Security - disable direct fallback on proxy error viper.SetDefault("security.proxy_fallback.allow_direct_on_error", false) @@ -1875,8 +2006,14 @@ func setDefaults() { viper.SetDefault("ops.enabled", true) viper.SetDefault("ops.use_preaggregated_tables", true) viper.SetDefault("ops.cleanup.enabled", true) - viper.SetDefault("ops.cleanup.schedule", "0 2 * * *") - viper.SetDefault("ops.cleanup.archive_expire_days", 14) + viper.SetDefault("ops.cleanup.schedule", "0 4 * * *") + viper.SetDefault("ops.cleanup.archive_expire_days", 30) + viper.SetDefault("ops.cleanup.archive_window_days", 1) + viper.SetDefault("ops.cleanup.max_catchup_windows_per_run", 2) + viper.SetDefault("ops.cleanup.archive_timeout_seconds", 1800) + viper.SetDefault("ops.cleanup.delete_timeout_seconds", 1800) + viper.SetDefault("ops.cleanup.run_timeout_seconds", 18000) + viper.SetDefault("ops.cleanup.delete_batch_size", 5000) // Retention days: vNext defaults to 30 days across ops datasets. viper.SetDefault("ops.cleanup.error_log_retention_days", 30) viper.SetDefault("ops.cleanup.minute_metrics_retention_days", 30) @@ -1992,8 +2129,10 @@ func setDefaults() { viper.SetDefault("gateway.response_header_timeout", 600) // 600秒(10分钟)等待上游响应头，LLM高负载时可能排队较久 viper.SetDefault("gateway.openai_response_header_timeout", 600) viper.SetDefault("gateway.image_nonstream_total_timeout_seconds", 1800) - viper.SetDefault("gateway.openai_first_output_timeout_seconds", 0) - viper.SetDefault("gateway.openai_high_effort_first_output_timeout_seconds", 0) + // 首输出保护针对 OpenAI HTTP Responses（含 passthrough）的语义事件；已开始输出后不再计时。 + // 60/180 秒覆盖正常请求，同时把坏代理/无响应上游的极端长尾转为一次受控 failover。 + viper.SetDefault("gateway.openai_first_output_timeout_seconds", 60) + viper.SetDefault("gateway.openai_high_effort_first_output_timeout_seconds", 180) viper.SetDefault("gateway.log_upstream_error_body", true) viper.SetDefault("gateway.log_upstream_error_body_max_bytes", 2048) viper.SetDefault("gateway.inject_beta_for_apikey", false) @@ -2147,6 +2286,21 @@ func setDefaults() { } func (c *Config) Validate() error { + c.Server.TrustedProxies = normalizeStringSlice(c.Server.TrustedProxies) + for _, trustedProxy := range c.Server.TrustedProxies { + if net.ParseIP(trustedProxy) != nil { + continue + } + if _, _, err := net.ParseCIDR(trustedProxy); err != nil { + return fmt.Errorf("server.trusted_proxies contains invalid IP or CIDR %q", trustedProxy) + } + } + forwardedClientIPHeaders, err := NormalizeForwardedClientIPHeaders(c.Security.ForwardedClientIPHeaders) + if err != nil { + return fmt.Errorf("security.forwarded_client_ip_headers: %w", err) + } + c.Security.ForwardedClientIPHeaders = forwardedClientIPHeaders + jwtSecret := strings.TrimSpace(c.JWT.Secret) if jwtSecret == "" { return fmt.Errorf("jwt.secret is required") @@ -2274,6 +2428,9 @@ func (c *Config) Validate() error { if _, err := c.Server.ListenSpec(); err != nil { return fmt.Errorf("server listen config invalid: %w", err) } + if c.Server.ShutdownTimeoutSeconds < 0 || c.Server.ShutdownTimeoutSeconds > maxServerShutdownTimeoutSeconds { + return fmt.Errorf("server.shutdown_timeout_seconds must be 0 or between 1-%d seconds", maxServerShutdownTimeoutSeconds) + } if c.JWT.ExpireHour <= 0 { return fmt.Errorf("jwt.expire_hour must be positive") } @@ -3050,20 +3207,8 @@ func (c *Config) Validate() error { if c.Ops.MetricsCollectorCache.TTL < 0 { return fmt.Errorf("ops.metrics_collector_cache.ttl must be non-negative") } - if c.Ops.Cleanup.ArchiveExpireDays < 0 { - return fmt.Errorf("ops.cleanup.archive_expire_days must be non-negative") - } - if c.Ops.Cleanup.ErrorLogRetentionDays < 0 { - return fmt.Errorf("ops.cleanup.error_log_retention_days must be non-negative") - } - if c.Ops.Cleanup.MinuteMetricsRetentionDays < 0 { - return fmt.Errorf("ops.cleanup.minute_metrics_retention_days must be non-negative") - } - if c.Ops.Cleanup.HourlyMetricsRetentionDays < 0 { - return fmt.Errorf("ops.cleanup.hourly_metrics_retention_days must be non-negative") - } - if c.Ops.Cleanup.Enabled && strings.TrimSpace(c.Ops.Cleanup.Schedule) == "" { - return fmt.Errorf("ops.cleanup.schedule is required when ops.cleanup.enabled=true") + if err := c.Ops.Cleanup.Validate(); err != nil { + return err } if c.Concurrency.PingInterval < 5 || c.Concurrency.PingInterval > 30 { return fmt.Errorf("concurrency.ping_interval must be between 5-30 seconds") diff --git a/backend/internal/config/config_test.go b/backend/internal/config/config_test.go index 69ed394eb..2c81f1551 100644 --- a/backend/internal/config/config_test.go +++ b/backend/internal/config/config_test.go @@ -1,6 +1,7 @@ package config import ( + "fmt" "net" "os" pathpkg "path" @@ -14,6 +15,66 @@ import ( "github.com/stretchr/testify/require" ) +func TestNormalizeForwardedClientIPHeaders(t *testing.T) { + headers, err := NormalizeForwardedClientIPHeaders([]string{ + " x-cdn-client-ip ", + "X-CDN-CLIENT-IP", + "true-client-ip", + }) + require.NoError(t, err) + require.Equal(t, []string{"X-Cdn-Client-Ip", "True-Client-Ip"}, headers) + + _, err = NormalizeForwardedClientIPHeaders([]string{"X Invalid"}) + require.ErrorContains(t, err, "invalid HTTP header field name") + + _, err = NormalizeForwardedClientIPHeaders([]string{"Authorization"}) + require.ErrorContains(t, err, "cannot be used as a client IP source") +} + +func TestNormalizeForwardedClientIPHeadersLimit(t *testing.T) { + headers := make([]string, 0, MaxForwardedClientIPHeaders+1) + for i := 0; i <= MaxForwardedClientIPHeaders; i++ { + headers = append(headers, fmt.Sprintf("X-CDN-IP-%d", i)) + } + + _, err := NormalizeForwardedClientIPHeaders(headers) + require.ErrorContains(t, err, "at most 16 unique names") +} + +func TestLoadForwardedClientIPHeaders(t *testing.T) { + t.Run("yaml normalizes and deduplicates", func(t *testing.T) { + resetViperWithConfig(t, "security:\n forwarded_client_ip_headers: [x-cdn-ip, X-CDN-IP, true-client-ip]\njwt:\n secret: "+strings.Repeat("x", 32)+"\ntotp:\n encryption_key: "+strings.Repeat("a", 64)+"\n") + + cfg, err := Load() + require.NoError(t, err) + require.Equal(t, []string{"X-Cdn-Ip", "True-Client-Ip"}, cfg.Security.ForwardedClientIPHeaders) + }) + + t.Run("environment overrides yaml", func(t *testing.T) { + resetViperWithJWTSecret(t) + t.Setenv("SECURITY_FORWARDED_CLIENT_IP_HEADERS", " x-env-ip , X-ENV-IP, true-client-ip ") + + cfg, err := Load() + require.NoError(t, err) + require.Equal(t, []string{"X-Env-Ip", "True-Client-Ip"}, cfg.Security.ForwardedClientIPHeaders) + }) + + t.Run("invalid header fails fast", func(t *testing.T) { + resetViperWithJWTSecret(t) + t.Setenv("SECURITY_FORWARDED_CLIENT_IP_HEADERS", "X-Valid-IP, X Invalid") + + _, err := Load() + require.ErrorContains(t, err, "security.forwarded_client_ip_headers") + }) +} + +func TestConfigValidateRejectsInvalidTrustedProxy(t *testing.T) { + resetViperWithConfig(t, "server:\n trusted_proxies: [not-a-cidr]\njwt:\n secret: "+strings.Repeat("x", 32)+"\ntotp:\n encryption_key: "+strings.Repeat("a", 64)+"\n") + + _, err := Load() + require.ErrorContains(t, err, "server.trusted_proxies contains invalid IP or CIDR") +} + func resetViperWithJWTSecret(t *testing.T) { t.Helper() resetViperWithEmptyConfig(t) @@ -60,6 +121,56 @@ func TestLoadServerTimingConfig(t *testing.T) { }) } +func TestLoadServerShutdownTimeout(t *testing.T) { + t.Run("default", func(t *testing.T) { + resetViperWithJWTSecret(t) + + cfg, err := Load() + require.NoError(t, err) + require.Equal(t, 30, cfg.Server.ShutdownTimeoutSeconds) + require.Equal(t, 30*time.Second, cfg.Server.ShutdownTimeout()) + }) + + t.Run("environment override", func(t *testing.T) { + resetViperWithJWTSecret(t) + t.Setenv("SERVER_SHUTDOWN_TIMEOUT_SECONDS", "45") + + cfg, err := Load() + require.NoError(t, err) + require.Equal(t, 45, cfg.Server.ShutdownTimeoutSeconds) + require.Equal(t, 45*time.Second, cfg.Server.ShutdownTimeout()) + }) + + t.Run("explicit zero uses safe fallback", func(t *testing.T) { + resetViperWithConfig(t, "server:\n shutdown_timeout_seconds: 0\njwt:\n secret: "+strings.Repeat("x", 32)+"\ntotp:\n encryption_key: "+strings.Repeat("a", 64)+"\n") + + cfg, err := Load() + require.NoError(t, err) + require.Zero(t, cfg.Server.ShutdownTimeoutSeconds) + require.Equal(t, 30*time.Second, cfg.Server.ShutdownTimeout()) + }) + + for _, testCase := range []struct { + name string + value string + }{ + {name: "negative", value: "-1"}, + {name: "over maximum", value: "3601"}, + } { + t.Run(testCase.name, func(t *testing.T) { + resetViperWithJWTSecret(t) + t.Setenv("SERVER_SHUTDOWN_TIMEOUT_SECONDS", testCase.value) + + _, err := Load() + require.ErrorContains(t, err, "server.shutdown_timeout_seconds") + }) + } +} + +func TestServerConfigShutdownTimeoutZeroValue(t *testing.T) { + require.Equal(t, 30*time.Second, (ServerConfig{}).ShutdownTimeout()) +} + func resetViperWithConfig(t *testing.T, content string) string { t.Helper() viper.Reset() @@ -150,6 +261,8 @@ func TestLoadDefaultOpenAIImageTimeoutConfig(t *testing.T) { cfg, err := Load() require.NoError(t, err) require.Equal(t, 600, cfg.Gateway.OpenAIResponseHeaderTimeout) + require.Equal(t, 60, cfg.Gateway.OpenAIFirstOutputTimeoutSeconds) + require.Equal(t, 180, cfg.Gateway.OpenAIHighEffortFirstOutputTimeoutSeconds) require.Equal(t, 1800, cfg.Gateway.ImageNonstreamTotalTimeoutSeconds) require.Equal(t, DefaultUpstreamResponseReadMaxBytes, cfg.Gateway.UpstreamResponseReadMaxBytes) } @@ -1161,6 +1274,27 @@ func TestValidateOpsCleanupScheduleRequired(t *testing.T) { } } +func TestOpsCleanupDefaults(t *testing.T) { + resetViperWithJWTSecret(t) + cfg, err := Load() + if err != nil { + t.Fatalf("Load() error: %v", err) + } + cleanup := cfg.Ops.Cleanup + if cleanup.Schedule != "0 4 * * *" { + t.Fatalf("schedule = %q, want 0 4 * * *", cleanup.Schedule) + } + if cleanup.ArchiveExpireDays != 30 { + t.Fatalf("archive_expire_days = %d, want 30", cleanup.ArchiveExpireDays) + } + if cleanup.ArchiveWindowDays != 1 || cleanup.MaxCatchupWindowsPerRun != 2 { + t.Fatalf("window defaults = %d/%d, want 1/2", cleanup.ArchiveWindowDays, cleanup.MaxCatchupWindowsPerRun) + } + if cleanup.ArchiveTimeoutSeconds <= 0 || cleanup.DeleteTimeoutSeconds <= 0 || cleanup.RunTimeoutSeconds <= 0 || cleanup.DeleteBatchSize <= 0 { + t.Fatalf("cleanup execution controls must be positive: %+v", cleanup) + } +} + func TestValidateConcurrencyPingInterval(t *testing.T) { resetViperWithJWTSecret(t) @@ -1750,6 +1884,26 @@ func TestValidateConfigErrors(t *testing.T) { mutate: func(c *Config) { c.Ops.Cleanup.MinuteMetricsRetentionDays = -1 }, wantErr: "ops.cleanup.minute_metrics_retention_days", }, + { + name: "ops cleanup archive window", + mutate: func(c *Config) { c.Ops.Cleanup.ArchiveWindowDays = 0 }, + wantErr: "ops.cleanup.archive_window_days", + }, + { + name: "ops cleanup archive timeout", + mutate: func(c *Config) { c.Ops.Cleanup.ArchiveTimeoutSeconds = 0 }, + wantErr: "ops.cleanup.archive_timeout_seconds", + }, + { + name: "ops cleanup delete timeout", + mutate: func(c *Config) { c.Ops.Cleanup.DeleteTimeoutSeconds = 0 }, + wantErr: "ops.cleanup.delete_timeout_seconds", + }, + { + name: "ops cleanup run timeout budget", + mutate: func(c *Config) { c.Ops.Cleanup.RunTimeoutSeconds = 1 }, + wantErr: "ops.cleanup.run_timeout_seconds", + }, } for _, tt := range cases { diff --git a/backend/internal/handler/admin/account_handler.go b/backend/internal/handler/admin/account_handler.go index 1962d9744..0d89cc4dd 100644 --- a/backend/internal/handler/admin/account_handler.go +++ b/backend/internal/handler/admin/account_handler.go @@ -948,6 +948,7 @@ func (h *AccountHandler) Update(c *gin.Context) { return } + h.enqueueOwnedPublicShareValidation(account) response.Success(c, h.buildAccountResponseWithRuntime(c.Request.Context(), account)) } diff --git a/backend/internal/handler/admin/account_handler_public_share_update_test.go b/backend/internal/handler/admin/account_handler_public_share_update_test.go new file mode 100644 index 000000000..d71d25863 --- /dev/null +++ b/backend/internal/handler/admin/account_handler_public_share_update_test.go @@ -0,0 +1,65 @@ +package admin + +import ( + "bytes" + "context" + "net/http" + "net/http/httptest" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +type pendingOwnedAccountUpdateAdminService struct { + *stubAdminService + updatedAccount *service.Account +} + +func (s *pendingOwnedAccountUpdateAdminService) UpdateAccount(context.Context, int64, *service.UpdateAccountInput) (*service.Account, error) { + return s.updatedAccount, nil +} + +func TestAccountHandlerUpdateEnqueuesOwnedPendingPublicShareValidation(t *testing.T) { + gin.SetMode(gin.TestMode) + ownerUserID := int64(101) + adminService := &pendingOwnedAccountUpdateAdminService{ + stubAdminService: newStubAdminService(), + updatedAccount: &service.Account{ + ID: 7, + Name: "owned-pending", + Platform: service.PlatformOpenAI, + Type: service.AccountTypeOAuth, + OwnerUserID: &ownerUserID, + ShareMode: service.AccountShareModePublic, + ShareStatus: service.AccountShareStatusPending, + Status: service.StatusActive, + }, + } + handler := &AccountHandler{ + adminService: adminService, + accountService: &service.AccountService{}, + accountTestService: &service.AccountTestService{}, + publicShareValidation: make(chan ownedPublicShareValidationJob, 1), + } + // Keep the test deterministic: mark worker startup complete so the queued + // job remains available for direct assertion instead of being consumed. + handler.publicShareValidationOnce.Do(func() {}) + router := gin.New() + router.PUT("/api/v1/admin/accounts/:id", handler.Update) + + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodPut, "/api/v1/admin/accounts/7", bytes.NewBufferString(`{}`)) + request.Header.Set("Content-Type", "application/json") + router.ServeHTTP(recorder, request) + + require.Equal(t, http.StatusOK, recorder.Code) + select { + case job := <-handler.publicShareValidation: + require.Equal(t, int64(7), job.AccountID) + require.Equal(t, ownerUserID, job.OwnerUserID) + default: + t.Fatal("expected pending owned account validation job") + } +} diff --git a/backend/internal/handler/admin/setting_handler.go b/backend/internal/handler/admin/setting_handler.go index 7c8b5480f..e63b8741c 100644 --- a/backend/internal/handler/admin/setting_handler.go +++ b/backend/internal/handler/admin/setting_handler.go @@ -234,6 +234,7 @@ func (h *SettingHandler) GetSettings(c *gin.Context) { WithdrawalManagementEnabled: settings.WithdrawalManagementEnabled, WithdrawalRateLimitWindowDays: settings.WithdrawalRateLimitWindowDays, WithdrawalRateLimitMax: settings.WithdrawalRateLimitMax, + WithdrawalRateLimitExemptAmount: settings.WithdrawalRateLimitExemptAmount, CyberSessionBlockEnabled: settings.CyberSessionBlockEnabled, CyberSessionBlockTTLSeconds: settings.CyberSessionBlockTTLSeconds, AccountShareCommentReviewEnabled: settings.AccountShareCommentReviewEnabled, @@ -710,10 +711,11 @@ type UpdateSettingsRequest struct { AvailableChannelsEnabled *bool `json:"available_channels_enabled"` // Functional module switches - InvoiceManagementEnabled *bool `json:"invoice_management_enabled"` - WithdrawalManagementEnabled *bool `json:"withdrawal_management_enabled"` - WithdrawalRateLimitWindowDays *int `json:"withdrawal_rate_limit_window_days"` - WithdrawalRateLimitMax *int `json:"withdrawal_rate_limit_max"` + InvoiceManagementEnabled *bool `json:"invoice_management_enabled"` + WithdrawalManagementEnabled *bool `json:"withdrawal_management_enabled"` + WithdrawalRateLimitWindowDays *int `json:"withdrawal_rate_limit_window_days"` + WithdrawalRateLimitMax *int `json:"withdrawal_rate_limit_max"` + WithdrawalRateLimitExemptAmount *float64 `json:"withdrawal_rate_limit_exempt_amount"` // User-owned account import limit UserAccountImportLimit *int `json:"user_account_import_limit"` @@ -795,6 +797,23 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) { return } } + if req.WithdrawalRateLimitExemptAmount != nil { + config := service.WithdrawalRateLimitConfig{ + WindowDays: previousSettings.WithdrawalRateLimitWindowDays, + MaxRequests: previousSettings.WithdrawalRateLimitMax, + ExemptAmount: *req.WithdrawalRateLimitExemptAmount, + } + if req.WithdrawalRateLimitWindowDays != nil { + config.WindowDays = *req.WithdrawalRateLimitWindowDays + } + if req.WithdrawalRateLimitMax != nil { + config.MaxRequests = *req.WithdrawalRateLimitMax + } + if err := service.ValidateWithdrawalRateLimitConfig(config); err != nil { + response.ErrorFrom(c, err) + return + } + } if req.UserAccountImportLimit != nil { value := service.NormalizeUserAccountCredentialImportLimit(*req.UserAccountImportLimit) req.UserAccountImportLimit = &value @@ -1607,6 +1626,12 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) { } return previousSettings.WithdrawalRateLimitMax }(), + WithdrawalRateLimitExemptAmount: func() float64 { + if req.WithdrawalRateLimitExemptAmount != nil { + return *req.WithdrawalRateLimitExemptAmount + } + return previousSettings.WithdrawalRateLimitExemptAmount + }(), CyberSessionBlockEnabled: func() bool { if req.CyberSessionBlockEnabled != nil { return *req.CyberSessionBlockEnabled @@ -2109,6 +2134,7 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) { WithdrawalManagementEnabled: updatedSettings.WithdrawalManagementEnabled, WithdrawalRateLimitWindowDays: updatedSettings.WithdrawalRateLimitWindowDays, WithdrawalRateLimitMax: updatedSettings.WithdrawalRateLimitMax, + WithdrawalRateLimitExemptAmount: updatedSettings.WithdrawalRateLimitExemptAmount, CyberSessionBlockEnabled: updatedSettings.CyberSessionBlockEnabled, CyberSessionBlockTTLSeconds: updatedSettings.CyberSessionBlockTTLSeconds, AccountShareCommentReviewEnabled: updatedSettings.AccountShareCommentReviewEnabled, @@ -2528,6 +2554,9 @@ func preserveOmittedUpdateSettingsFields(req *UpdateSettingsRequest, previous *s if !fieldProvided(fields, "withdrawal_rate_limit_max") { req.WithdrawalRateLimitMax = &previous.WithdrawalRateLimitMax } + if !fieldProvided(fields, "withdrawal_rate_limit_exempt_amount") { + req.WithdrawalRateLimitExemptAmount = &previous.WithdrawalRateLimitExemptAmount + } if !fieldProvided(fields, "cyber_session_block_enabled") { req.CyberSessionBlockEnabled = &previous.CyberSessionBlockEnabled } @@ -3057,6 +3086,9 @@ func diffSettings(before *service.SystemSettings, after *service.SystemSettings, if before.WithdrawalRateLimitMax != after.WithdrawalRateLimitMax { changed = append(changed, "withdrawal_rate_limit_max") } + if before.WithdrawalRateLimitExemptAmount != after.WithdrawalRateLimitExemptAmount { + changed = append(changed, "withdrawal_rate_limit_exempt_amount") + } if before.UserAccountImportLimit != after.UserAccountImportLimit { changed = append(changed, "user_account_import_limit") } diff --git a/backend/internal/handler/api_key_handler.go b/backend/internal/handler/api_key_handler.go index d5d1fbaea..357cbf023 100644 --- a/backend/internal/handler/api_key_handler.go +++ b/backend/internal/handler/api_key_handler.go @@ -59,6 +59,7 @@ type CreateAPIKeyRequest struct { IPBlacklist []string `json:"ip_blacklist"` // IP 黑名单 Quota *float64 `json:"quota"` // 配额限制 (USD) ExpiresInDays *int `json:"expires_in_days"` // 过期天数 + ExpiresAt *string `json:"expires_at"` // 精确过期时间 (RFC3339) // Rate limit fields (0 = unlimited) RateLimit5h *float64 `json:"rate_limit_5h"` @@ -72,8 +73,8 @@ type UpdateAPIKeyRequest struct { GroupID *int64 `json:"group_id"` GroupRoutes *[]APIKeyGroupRouteRequest `json:"group_routes"` Status string `json:"status" binding:"omitempty,oneof=active inactive"` - IPWhitelist []string `json:"ip_whitelist"` // IP 白名单 - IPBlacklist []string `json:"ip_blacklist"` // IP 黑名单 + IPWhitelist *[]string `json:"ip_whitelist"` // IP 白名单（nil 不修改，空数组清空） + IPBlacklist *[]string `json:"ip_blacklist"` // IP 黑名单（nil 不修改，空数组清空） Quota *float64 `json:"quota"` // 配额限制 (USD), 0=无限制 ExpiresAt *string `json:"expires_at"` // 过期时间 (ISO 8601) ResetQuota *bool `json:"reset_quota"` // 重置已用配额 @@ -93,6 +94,23 @@ type APIKeyGroupRouteRequest struct { CooldownSeconds int `json:"cooldown_seconds"` } +func parseCreateAPIKeyExpiration(rawExpiresAt *string, expiresInDays *int, now time.Time) (*time.Time, error) { + if rawExpiresAt != nil && expiresInDays != nil { + return nil, service.ErrAPIKeyExpirationConflict + } + if rawExpiresAt == nil { + return nil, nil + } + expiresAt, err := time.Parse(time.RFC3339, *rawExpiresAt) + if err != nil { + return nil, service.ErrAPIKeyExpirationInvalid.WithCause(err) + } + if !expiresAt.After(now) { + return nil, service.ErrAPIKeyExpirationNotFuture + } + return &expiresAt, nil +} + // List handles listing user's API keys with pagination // GET /api/v1/api-keys func (h *APIKeyHandler) List(c *gin.Context) { @@ -183,6 +201,11 @@ func (h *APIKeyHandler) Create(c *gin.Context) { response.BadRequest(c, "Invalid request: "+err.Error()) return } + expiresAt, err := parseCreateAPIKeyExpiration(req.ExpiresAt, req.ExpiresInDays, time.Now()) + if err != nil { + response.ErrorFrom(c, err) + return + } svcReq := service.CreateAPIKeyRequest{ Name: req.Name, @@ -192,6 +215,7 @@ func (h *APIKeyHandler) Create(c *gin.Context) { IPWhitelist: req.IPWhitelist, IPBlacklist: req.IPBlacklist, ExpiresInDays: req.ExpiresInDays, + ExpiresAt: expiresAt, } if req.Quota != nil { svcReq.Quota = *req.Quota diff --git a/backend/internal/handler/api_key_handler_create_expiration_test.go b/backend/internal/handler/api_key_handler_create_expiration_test.go new file mode 100644 index 000000000..dacc11903 --- /dev/null +++ b/backend/internal/handler/api_key_handler_create_expiration_test.go @@ -0,0 +1,64 @@ +package handler + +import ( + "bytes" + "net/http" + "net/http/httptest" + "testing" + "time" + + middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestParseCreateAPIKeyExpirationPreservesRFC3339Instant(t *testing.T) { + now := time.Date(2026, time.July, 23, 8, 0, 0, 0, time.UTC) + raw := "2026-07-24T16:17:18+08:00" + + expiresAt, err := parseCreateAPIKeyExpiration(&raw, nil, now) + + require.NoError(t, err) + require.NotNil(t, expiresAt) + require.Equal(t, raw, expiresAt.Format(time.RFC3339)) +} + +func TestAPIKeyHandlerCreateRejectsInvalidExpiration(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + h := NewAPIKeyHandler(nil) + router.POST("/keys", func(c *gin.Context) { + c.Set(string(middleware2.ContextKeyUser), middleware2.AuthSubject{UserID: 42}) + c.Next() + }, h.Create) + + tests := []struct { + name string + body string + }{ + { + name: "expires_at and expires_in_days conflict", + body: `{"name":"conflict","expires_at":"2099-03-03T21:06:00Z","expires_in_days":30}`, + }, + { + name: "expires_at is not RFC3339", + body: `{"name":"invalid format","expires_at":"2099-03-03 21:06:00"}`, + }, + { + name: "expires_at is in the past", + body: `{"name":"past","expires_at":"2000-01-01T00:00:00Z"}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodPost, "/keys", bytes.NewBufferString(tt.body)) + request.Header.Set("Content-Type", "application/json") + + router.ServeHTTP(recorder, request) + + require.Equal(t, http.StatusBadRequest, recorder.Code, recorder.Body.String()) + }) + } +} diff --git a/backend/internal/handler/auth_handler.go b/backend/internal/handler/auth_handler.go index bd85dfe77..ff966a100 100644 --- a/backend/internal/handler/auth_handler.go +++ b/backend/internal/handler/auth_handler.go @@ -174,7 +174,7 @@ func (h *AuthHandler) Register(c *gin.Context) { } // Turnstile 验证（邮箱验证码注册场景避免重复校验一次性 token） - if err := h.authService.VerifyTurnstileForRegister(c.Request.Context(), req.TurnstileToken, ip.GetClientIP(c), req.VerifyCode); err != nil { + if err := h.authService.VerifyTurnstileForRegister(c.Request.Context(), req.TurnstileToken, ip.GetSecurityClientIP(c), req.VerifyCode); err != nil { response.ErrorFrom(c, err) return } @@ -210,7 +210,7 @@ func (h *AuthHandler) SendVerifyCode(c *gin.Context) { } // Turnstile 验证 - if err := h.authService.VerifyTurnstile(c.Request.Context(), req.TurnstileToken, ip.GetClientIP(c)); err != nil { + if err := h.authService.VerifyTurnstile(c.Request.Context(), req.TurnstileToken, ip.GetSecurityClientIP(c)); err != nil { response.ErrorFrom(c, err) return } @@ -237,7 +237,7 @@ func (h *AuthHandler) Login(c *gin.Context) { } // Turnstile 验证 - if err := h.authService.VerifyTurnstile(c.Request.Context(), req.TurnstileToken, ip.GetClientIP(c)); err != nil { + if err := h.authService.VerifyTurnstile(c.Request.Context(), req.TurnstileToken, ip.GetSecurityClientIP(c)); err != nil { response.ErrorFrom(c, err) return } @@ -625,7 +625,7 @@ func (h *AuthHandler) ForgotPassword(c *gin.Context) { } // Turnstile 验证 - if err := h.authService.VerifyTurnstile(c.Request.Context(), req.TurnstileToken, ip.GetClientIP(c)); err != nil { + if err := h.authService.VerifyTurnstile(c.Request.Context(), req.TurnstileToken, ip.GetSecurityClientIP(c)); err != nil { response.ErrorFrom(c, err) return } diff --git a/backend/internal/handler/auth_oauth_pending_flow.go b/backend/internal/handler/auth_oauth_pending_flow.go index c8e278c11..4b6cb0d5f 100644 --- a/backend/internal/handler/auth_oauth_pending_flow.go +++ b/backend/internal/handler/auth_oauth_pending_flow.go @@ -602,7 +602,7 @@ func (h *AuthHandler) SendPendingOAuthVerifyCode(c *gin.Context) { return } - if err := h.authService.VerifyTurnstile(c.Request.Context(), req.TurnstileToken, ip.GetClientIP(c)); err != nil { + if err := h.authService.VerifyTurnstile(c.Request.Context(), req.TurnstileToken, ip.GetSecurityClientIP(c)); err != nil { response.ErrorFrom(c, err) return } diff --git a/backend/internal/handler/dto/mappers.go b/backend/internal/handler/dto/mappers.go index 3dddb28fc..59d8dabfb 100644 --- a/backend/internal/handler/dto/mappers.go +++ b/backend/internal/handler/dto/mappers.go @@ -672,6 +672,8 @@ func usageLogFromServiceUser(l *service.UsageLog) UsageLog { CacheReadTokens: l.CacheReadTokens, CacheCreation5mTokens: l.CacheCreation5mTokens, CacheCreation1hTokens: l.CacheCreation1hTokens, + ImageInputTokens: l.ImageInputTokens, + ImageInputCost: l.ImageInputCost, InputCost: l.InputCost, OutputCost: l.OutputCost, CacheCreationCost: l.CacheCreationCost, diff --git a/backend/internal/handler/dto/mappers_usage_test.go b/backend/internal/handler/dto/mappers_usage_test.go index b5b37247e..7e6c05ce0 100644 --- a/backend/internal/handler/dto/mappers_usage_test.go +++ b/backend/internal/handler/dto/mappers_usage_test.go @@ -150,6 +150,27 @@ func TestUsageLogFromService_FallsBackToLegacyModelWhenRequestedModelMissing(t * require.Equal(t, "claude-3", adminDTO.Model) } +func TestUsageLogFromService_IncludesImageInputUsage(t *testing.T) { + t.Parallel() + + log := &service.UsageLog{ + RequestID: "req_image_input", + Model: "gpt-image-2", + InputTokens: 371, + ImageInputTokens: 352, + InputCost: 0.000095, + ImageInputCost: 0.002816, + } + + userDTO := UsageLogFromService(log) + adminDTO := UsageLogFromServiceAdmin(log) + + require.Equal(t, 352, userDTO.ImageInputTokens) + require.InDelta(t, 0.002816, userDTO.ImageInputCost, 1e-15) + require.Equal(t, 352, adminDTO.ImageInputTokens) + require.InDelta(t, 0.002816, adminDTO.ImageInputCost, 1e-15) +} + func f64Ptr(value float64) *float64 { return &value } diff --git a/backend/internal/handler/dto/settings.go b/backend/internal/handler/dto/settings.go index 3dac3abb6..fd60ea114 100644 --- a/backend/internal/handler/dto/settings.go +++ b/backend/internal/handler/dto/settings.go @@ -273,10 +273,11 @@ type SystemSettings struct { AvailableChannelsEnabled bool `json:"available_channels_enabled"` // Functional module switches - InvoiceManagementEnabled bool `json:"invoice_management_enabled"` - WithdrawalManagementEnabled bool `json:"withdrawal_management_enabled"` - WithdrawalRateLimitWindowDays int `json:"withdrawal_rate_limit_window_days"` - WithdrawalRateLimitMax int `json:"withdrawal_rate_limit_max"` + InvoiceManagementEnabled bool `json:"invoice_management_enabled"` + WithdrawalManagementEnabled bool `json:"withdrawal_management_enabled"` + WithdrawalRateLimitWindowDays int `json:"withdrawal_rate_limit_window_days"` + WithdrawalRateLimitMax int `json:"withdrawal_rate_limit_max"` + WithdrawalRateLimitExemptAmount float64 `json:"withdrawal_rate_limit_exempt_amount"` // User-owned account import limit UserAccountImportLimit int `json:"user_account_import_limit"` @@ -350,12 +351,13 @@ type PublicSettings struct { OpenAIAccountLevels []OpenAIAccountLevelConfig `json:"openai_account_levels"` - AffiliateEnabled bool `json:"affiliate_enabled"` - InvoiceManagementEnabled bool `json:"invoice_management_enabled"` - WithdrawalManagementEnabled bool `json:"withdrawal_management_enabled"` - WithdrawalRateLimitWindowDays int `json:"withdrawal_rate_limit_window_days"` - WithdrawalRateLimitMax int `json:"withdrawal_rate_limit_max"` - RiskControlEnabled bool `json:"risk_control_enabled"` + AffiliateEnabled bool `json:"affiliate_enabled"` + InvoiceManagementEnabled bool `json:"invoice_management_enabled"` + WithdrawalManagementEnabled bool `json:"withdrawal_management_enabled"` + WithdrawalRateLimitWindowDays int `json:"withdrawal_rate_limit_window_days"` + WithdrawalRateLimitMax int `json:"withdrawal_rate_limit_max"` + WithdrawalRateLimitExemptAmount float64 `json:"withdrawal_rate_limit_exempt_amount"` + RiskControlEnabled bool `json:"risk_control_enabled"` } type LoginAgreementDocument struct { diff --git a/backend/internal/handler/dto/types.go b/backend/internal/handler/dto/types.go index 2d2d7f7dd..96b0ddbb1 100644 --- a/backend/internal/handler/dto/types.go +++ b/backend/internal/handler/dto/types.go @@ -440,8 +440,10 @@ type UsageLog struct { CacheCreationTokens int `json:"cache_creation_tokens"` CacheReadTokens int `json:"cache_read_tokens"` - CacheCreation5mTokens int `json:"cache_creation_5m_tokens"` - CacheCreation1hTokens int `json:"cache_creation_1h_tokens"` + CacheCreation5mTokens int `json:"cache_creation_5m_tokens"` + CacheCreation1hTokens int `json:"cache_creation_1h_tokens"` + ImageInputTokens int `json:"image_input_tokens"` + ImageInputCost float64 `json:"image_input_cost"` InputCost float64 `json:"input_cost"` OutputCost float64 `json:"output_cost"` diff --git a/backend/internal/handler/failover_loop.go b/backend/internal/handler/failover_loop.go index b323dd852..c6be80c3b 100644 --- a/backend/internal/handler/failover_loop.go +++ b/backend/internal/handler/failover_loop.go @@ -103,16 +103,17 @@ func (s *FailoverState) HandleFailoverErrorWithRetryLimit( return FailoverExhausted } - // 缓存计费判断 - if needForceCacheBilling(s.hasBoundSession, failoverErr) { + if retryLimit < 0 { + retryLimit = 0 + } + sameAccountRetry := failoverErr.RetryableOnSameAccount && s.SameAccountRetryCount[accountID] < retryLimit + // 粘性会话只有在实际切换账号时才强制缓存计费；同账号重试不能重复计费。 + if needForceCacheBilling(s.hasBoundSession, failoverErr, sameAccountRetry) { s.ForceCacheBilling = true } // 同账号重试：对 RetryableOnSameAccount 的临时性错误，先在同一账号上重试 - if retryLimit < 0 { - retryLimit = 0 - } - if failoverErr.RetryableOnSameAccount && s.SameAccountRetryCount[accountID] < retryLimit { + if sameAccountRetry { s.SameAccountRetryCount[accountID]++ logger.FromContext(ctx).Warn("gateway.failover_same_account_retry", zap.Int64("account_id", accountID), @@ -195,9 +196,9 @@ func (s *FailoverState) HandleSelectionExhausted(ctx context.Context) FailoverAc } // needForceCacheBilling 判断 failover 时是否需要强制缓存计费。 -// 粘性会话切换账号、或上游明确标记时，将 input_tokens 转为 cache_read 计费。 -func needForceCacheBilling(hasBoundSession bool, failoverErr *service.UpstreamFailoverError) bool { - return hasBoundSession || (failoverErr != nil && failoverErr.ForceCacheBilling) +// 粘性会话实际切换账号、或上游明确标记时，将 input_tokens 转为 cache_read 计费。 +func needForceCacheBilling(hasBoundSession bool, failoverErr *service.UpstreamFailoverError, sameAccountRetry bool) bool { + return (hasBoundSession && !sameAccountRetry) || (failoverErr != nil && failoverErr.ForceCacheBilling) } // failoverClientGone 判断下游客户端是否已断开。上游请求可能使用分离后的 diff --git a/backend/internal/handler/failover_loop_test.go b/backend/internal/handler/failover_loop_test.go index 2d51e156e..8965f36c4 100644 --- a/backend/internal/handler/failover_loop_test.go +++ b/backend/internal/handler/failover_loop_test.go @@ -227,7 +227,7 @@ func TestHandleFailoverError_BasicSwitch(t *testing.T) { // --------------------------------------------------------------------------- func TestHandleFailoverError_CacheBilling(t *testing.T) { - t.Run("hasBoundSession为true时设置ForceCacheBilling", func(t *testing.T) { + t.Run("hasBoundSession为true且实际切换时设置ForceCacheBilling", func(t *testing.T) { mock := &mockTempUnscheduler{} fs := NewFailoverState(3, true) // hasBoundSession=true err := newTestFailoverErr(500, false, false) @@ -236,6 +236,29 @@ func TestHandleFailoverError_CacheBilling(t *testing.T) { require.True(t, fs.ForceCacheBilling) }) + t.Run("同账号重试时仅凭hasBoundSession不设置ForceCacheBilling", func(t *testing.T) { + mock := &mockTempUnscheduler{} + fs := NewFailoverState(3, true) + err := newTestFailoverErr(400, true, false) + + action := fs.HandleFailoverErrorWithRetryLimit(context.Background(), mock, 100, "openai", 1, err) + + require.Equal(t, FailoverContinue, action) + require.False(t, fs.ForceCacheBilling) + require.Zero(t, fs.SwitchCount) + }) + + t.Run("同账号重试耗尽并实际切换时设置ForceCacheBilling", func(t *testing.T) { + mock := &mockTempUnscheduler{} + fs := NewFailoverState(3, true) + err := newTestFailoverErr(400, true, false) + + fs.HandleFailoverErrorWithRetryLimit(context.Background(), mock, 100, "openai", 0, err) + + require.True(t, fs.ForceCacheBilling) + require.Equal(t, 1, fs.SwitchCount) + }) + t.Run("failoverErr.ForceCacheBilling为true时设置", func(t *testing.T) { mock := &mockTempUnscheduler{} fs := NewFailoverState(3, false) @@ -549,13 +572,14 @@ func TestHandleFailoverError_IntegrationScenario(t *testing.T) { action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", retryErr) require.Equal(t, FailoverContinue, action) } - require.True(t, fs.ForceCacheBilling, "hasBoundSession=true 应设置 ForceCacheBilling") + require.False(t, fs.ForceCacheBilling, "同账号重试期间不应重复计入缓存费用") // 2. 账号 100 超过重试上限 → TempUnschedule + 切换 action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", retryErr) require.Equal(t, FailoverContinue, action) require.Equal(t, 1, fs.SwitchCount) require.Len(t, mock.calls, 1) + require.True(t, fs.ForceCacheBilling, "同账号重试耗尽并实际换号时应设置 ForceCacheBilling") // 3. 账号 200 遇到不可重试错误 → 直接切换 switchErr := newTestFailoverErr(500, false, false) diff --git a/backend/internal/handler/gateway_helper.go b/backend/internal/handler/gateway_helper.go index 37b84fb9d..80b8ce469 100644 --- a/backend/internal/handler/gateway_helper.go +++ b/backend/internal/handler/gateway_helper.go @@ -357,6 +357,11 @@ func (h *ConcurrencyHelper) waitForSlotWithPing(c *gin.Context, slotType string, // waitForSlotWithPingTimeout waits for a concurrency slot with a custom timeout. func (h *ConcurrencyHelper) waitForSlotWithPingTimeout(c *gin.Context, slotType string, id int64, maxConcurrency int, timeout time.Duration, isStream bool, streamStarted *bool, tryImmediate bool) (func(), error) { + remaining, budgetEnabled := service.OpenAIFirstOutputBudgetRemaining(c.Request.Context()) + if budgetEnabled && remaining <= 0 { + return nil, fmt.Errorf("%w: slot_type=%s", service.ErrOpenAIFirstOutputRoutingBudgetExceeded, slotType) + } + timeout = service.CapOpenAIFirstOutputWait(c.Request.Context(), timeout) ctx, cancel := context.WithTimeout(c.Request.Context(), timeout) defer cancel() @@ -404,6 +409,9 @@ func (h *ConcurrencyHelper) waitForSlotWithPingTimeout(c *gin.Context, slotType for { select { case <-ctx.Done(): + if remaining, enabled := service.OpenAIFirstOutputBudgetRemaining(c.Request.Context()); enabled && remaining <= 0 { + return nil, fmt.Errorf("%w: slot_type=%s", service.ErrOpenAIFirstOutputRoutingBudgetExceeded, slotType) + } return nil, &ConcurrencyError{ SlotType: slotType, IsTimeout: true, diff --git a/backend/internal/handler/grok_media.go b/backend/internal/handler/grok_media.go index aa0af0eff..0d73c2892 100644 --- a/backend/internal/handler/grok_media.go +++ b/backend/internal/handler/grok_media.go @@ -47,6 +47,11 @@ func (h *OpenAIGatewayHandler) GrokVideoStatus(c *gin.Context) { h.handleGrokMedia(c, service.GrokMediaEndpointVideoStatus, c.Param("request_id")) } +// GrokVideoContent proxies downloadable video content through the task's upstream account. +func (h *OpenAIGatewayHandler) GrokVideoContent(c *gin.Context) { + h.handleGrokMedia(c, service.GrokMediaEndpointVideoContent, c.Param("request_id")) +} + func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.GrokMediaEndpoint, requestID string) { streamStarted := false defer h.recoverResponsesPanic(c, &streamStarted) @@ -100,7 +105,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "model is required") return } - if endpoint == service.GrokMediaEndpointVideoStatus && strings.TrimSpace(requestID) == "" { + if endpoint.IsVideoLookupRequest() && strings.TrimSpace(requestID) == "" { h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "request_id is required") return } @@ -129,6 +134,25 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. requestCtx := context.WithValue(c.Request.Context(), ctxkey.ForcePlatform, service.PlatformGrok) c.Request = c.Request.WithContext(requestCtx) + sessionSeed := body + if len(sessionSeed) == 0 && strings.TrimSpace(requestID) != "" { + sessionSeed = []byte(requestID) + } + sessionHash := h.gatewayService.GenerateExplicitSessionHash(c, sessionSeed) + boundLookupAccountID := int64(0) + if endpoint.IsVideoLookupRequest() { + sessionHash = service.GrokMediaVideoRequestSessionHash(requestID, subject.UserID, apiKey.ID) + boundLookupAccountID, err = h.gatewayService.ResolveGrokMediaVideoRequestAccount( + requestCtx, apiKey.GroupID, requestID, subject.UserID, apiKey.ID, + ) + if err != nil || boundLookupAccountID <= 0 { + reqLog.Info("grok_media.video_lookup_owner_binding_missing", zap.Error(err)) + h.errorResponse(c, http.StatusNotFound, "not_found_error", "Video request not found") + return + } + } else { + sessionHash = service.GrokMediaSessionHash(sessionHash) + } subscription, _ := middleware2.GetSubscriptionFromContext(c) service.SetOpsLatencyMs(c, service.OpsAuthLatencyMsKey, time.Since(requestStart).Milliseconds()) @@ -151,24 +175,17 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. return } - sessionSeed := body - if len(sessionSeed) == 0 && strings.TrimSpace(requestID) != "" { - sessionSeed = []byte(requestID) - } - sessionHash := h.gatewayService.GenerateExplicitSessionHash(c, sessionSeed) - sessionHash = service.GrokMediaSessionHash(sessionHash) - if endpoint == service.GrokMediaEndpointVideoStatus { - sessionHash = service.GrokMediaVideoRequestSessionHash(requestID) - } failedAccountIDs := make(map[int64]struct{}) sameAccountRetryCount := make(map[int64]int) var lastFailoverErr *service.UpstreamFailoverError + mediaEligibilityRejected := false switchCount := 0 maxAccountSwitches := h.maxAccountSwitches if maxAccountSwitches <= 0 { maxAccountSwitches = 3 } routingStart := time.Now() + requiredCapability := grokMediaRequiredCapability(endpoint) for { if failoverClientGone(c) { @@ -180,6 +197,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. sessionHash, requestModel, failedAccountIDs, + requiredCapability, ) if err != nil { if failoverClientGone(c) { @@ -189,6 +207,16 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. zap.Error(err), zap.Int("excluded_account_count", len(failedAccountIDs)), ) + if boundLookupAccountID > 0 { + h.errorResponse(c, http.StatusNotFound, "not_found_error", "Video request not found") + return + } + if endpoint.IsGenerationRequest() && errors.Is(err, service.ErrNoAvailableAccounts) && + (len(failedAccountIDs) == 0 || (mediaEligibilityRejected && lastFailoverErr == nil)) { + markOpsRoutingCapacityLimited(c) + h.errorResponse(c, http.StatusServiceUnavailable, "grok_media_no_eligible_account", "No eligible Grok media accounts") + return + } if len(failedAccountIDs) == 0 { cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, requestModel, requestModel, service.PlatformGrok) h.errorResponse(c, cls.Status, cls.ErrType, cls.Message) @@ -202,10 +230,62 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. return } if selection == nil || selection.Account == nil { + if endpoint.IsGenerationRequest() { + markOpsRoutingCapacityLimited(c) + h.errorResponse(c, http.StatusServiceUnavailable, "grok_media_no_eligible_account", "No eligible Grok media accounts") + return + } + if boundLookupAccountID > 0 { + h.errorResponse(c, http.StatusNotFound, "not_found_error", "Video request not found") + return + } cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, requestModel, requestModel, service.PlatformGrok) h.errorResponse(c, cls.Status, cls.ErrType, cls.Message) return } + if boundLookupAccountID > 0 && selection.Account.ID != boundLookupAccountID { + if selection.Acquired && selection.ReleaseFunc != nil { + selection.ReleaseFunc() + } + if bindErr := h.gatewayService.BindGrokMediaVideoRequestAccount( + requestCtx, + apiKey.GroupID, + requestID, + subject.UserID, + apiKey.ID, + boundLookupAccountID, + ); bindErr != nil { + reqLog.Warn("grok_media.video_lookup_binding_restore_failed", zap.Error(bindErr)) + } + reqLog.Warn("grok_media.video_lookup_bound_account_mismatch", + zap.Int64("bound_account_id", boundLookupAccountID), + zap.Int64("selected_account_id", selection.Account.ID), + ) + h.errorResponse(c, http.StatusNotFound, "not_found_error", "Video request not found") + return + } + + account := selection.Account + if endpoint.IsGenerationRequest() { + eligible, eligibilityReason, eligibilityErr := h.ensureGrokMediaAccountEligibility(requestCtx, account) + if !eligible { + releaseRejectedGrokMediaSelection(selection) + mediaEligibilityRejected = true + failedAccountIDs[account.ID] = struct{}{} + reqLog.Warn("grok_media.account_eligibility_rejected", + zap.Int64("account_id", account.ID), + zap.String("reason", eligibilityReason), + zap.Bool("probe_failed", eligibilityErr != nil), + ) + if switchCount >= maxAccountSwitches { + markOpsRoutingCapacityLimited(c) + h.errorResponse(c, http.StatusServiceUnavailable, "grok_media_no_eligible_account", "No eligible Grok media accounts") + return + } + switchCount++ + continue + } + } reqLog.Debug("grok_media.account_schedule_decision", zap.String("layer", scheduleDecision.Layer), @@ -216,14 +296,19 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. zap.Float64("load_skew", scheduleDecision.LoadSkew), ) - account := selection.Account sessionHash = ensureOpenAIPoolModeSessionHash(sessionHash, account) setOpsSelectedAccount(c, account.ID, account.Platform) - accountReleaseFunc, accountAcquired := h.acquireResponsesAccountSlot(c, apiKey.GroupID, sessionHash, selection, false, &streamStarted, reqLog) + freshAccount, accountReleaseFunc, accountAcquired := h.acquireResponsesAccountSlot(c, requestCtx, apiKey.GroupID, sessionHash, service.OpenAIAccountDispatchRequirements{ + RequestedModel: requestModel, + RequiredTransport: service.OpenAIUpstreamTransportHTTPSSE, + RequiredEndpointCapability: requiredCapability, + RequiredPlatform: service.PlatformGrok, + }, selection, false, &streamStarted, reqLog) if !accountAcquired { return } + account = freshAccount service.SetOpsLatencyMs(c, service.OpsRoutingLatencyMsKey, time.Since(routingStart).Milliseconds()) forwardStart := time.Now() @@ -281,6 +366,10 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. continue } } + if endpoint.IsVideoLookupRequest() { + h.handleFailoverExhausted(c, failoverErr, false) + return + } h.gatewayService.RecordOpenAIAccountSwitch() failedAccountIDs[account.ID] = struct{}{} lastFailoverErr = failoverErr @@ -310,7 +399,9 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, nil) if endpoint.IsVideoMutationRequest() && strings.TrimSpace(result.ResponseID) != "" { - if err := h.gatewayService.BindGrokMediaVideoRequestAccount(requestCtx, apiKey.GroupID, result.ResponseID, account.ID); err != nil { + if err := h.gatewayService.BindGrokMediaVideoRequestAccount( + requestCtx, apiKey.GroupID, result.ResponseID, subject.UserID, apiKey.ID, account.ID, + ); err != nil { reqLog.Warn("grok_media.bind_video_request_account_failed", zap.Int64("account_id", account.ID), zap.String("request_id", result.ResponseID), @@ -329,6 +420,36 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. } } +func (h *OpenAIGatewayHandler) ensureGrokMediaAccountEligibility(ctx context.Context, account *service.Account) (bool, string, error) { + if account == nil { + return false, "missing_account", errors.New("grok media account is required") + } + eligible, reason := account.GrokMediaGenerationEligibility() + if eligible || reason != "billing_unobserved" { + return eligible, reason, nil + } + if h == nil || h.grokMediaEligibilityProber == nil { + return false, "billing_probe_unavailable", errors.New("grok media eligibility probe is not configured") + } + return h.grokMediaEligibilityProber.ProbeMediaEligibility(ctx, account.ID) +} + +func grokMediaRequiredCapability(endpoint service.GrokMediaEndpoint) service.OpenAIEndpointCapability { + if endpoint.IsGenerationRequest() { + return service.OpenAIEndpointCapabilityGrokMediaGeneration + } + return "" +} + +func releaseRejectedGrokMediaSelection(selection *service.AccountSelectionResult) { + if selection != nil && selection.Acquired && selection.ReleaseFunc != nil { + release := selection.ReleaseFunc + selection.Acquired = false + selection.ReleaseFunc = nil + release() + } +} + func shouldRecordGrokMediaUsage(endpoint service.GrokMediaEndpoint, requestModel string) bool { return endpoint.IsGenerationRequest() && strings.TrimSpace(requestModel) != "" } diff --git a/backend/internal/handler/openai_account_share_mode.go b/backend/internal/handler/openai_account_share_mode.go index 9579548b4..c71c6a366 100644 --- a/backend/internal/handler/openai_account_share_mode.go +++ b/backend/internal/handler/openai_account_share_mode.go @@ -45,6 +45,17 @@ func openAICompatibleRequestContext(ctx context.Context, apiKey *service.APIKey) return context.WithValue(ctx, ctxkey.ForcePlatform, service.PlatformGrok) } +// openAIResponsesDispatchContext removes the routing-only deadline before the +// upstream attempt while retaining request cancellation and route-scoped values. +func openAIResponsesDispatchContext(c *gin.Context, routingCtx context.Context, apiKey *service.APIKey) context.Context { + ctx := context.Background() + if c != nil && c.Request != nil { + ctx = c.Request.Context() + } + ctx = service.WithAccountShareModeRequestFromContext(ctx, routingCtx) + return openAICompatibleRequestContext(ctx, apiKey) +} + func (h *OpenAIGatewayHandler) handleAccountShareModeSelectionError(c *gin.Context, err error, streamStarted bool) bool { switch { case errors.Is(err, service.ErrAccountShareModeGroupUnbound): diff --git a/backend/internal/handler/openai_alpha_search.go b/backend/internal/handler/openai_alpha_search.go index b8e335d63..f8ad68f9d 100644 --- a/backend/internal/handler/openai_alpha_search.go +++ b/backend/internal/handler/openai_alpha_search.go @@ -183,10 +183,14 @@ func (h *OpenAIGatewayHandler) AlphaSearch(c *gin.Context) { account := selection.Account setOpsSelectedAccount(c, account.ID, account.Platform) - accountRelease, acquired := h.acquireResponsesAccountSlot(c, currentAPIKey.GroupID, sessionHash, selection, false, &streamStarted, reqLog) + freshAccount, accountRelease, acquired := h.acquireResponsesAccountSlot(c, selectionCtx, currentAPIKey.GroupID, sessionHash, service.OpenAIAccountDispatchRequirements{ + RequestedModel: selectionModel, + RequiredTransport: service.OpenAIUpstreamTransportHTTPSSE, + }, selection, false, &streamStarted, reqLog) if !acquired { return } + account = freshAccount service.SetOpsLatencyMs(c, service.OpsRoutingLatencyMsKey, time.Since(routingStart).Milliseconds()) writerSizeBeforeForward := c.Writer.Size() forwardBody := body @@ -203,7 +207,7 @@ func (h *OpenAIGatewayHandler) AlphaSearch(c *gin.Context) { service.SetOpsLatencyMs(c, service.OpsResponseLatencyMsKey, time.Since(forwardStart).Milliseconds()) if forwardErr == nil { - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, nil) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, nil, account.GetMappedModel(selectionModel)) routeCursor.recordSuccess(currentAPIKey.ID) if result != nil { userAgent := c.GetHeader("User-Agent") diff --git a/backend/internal/handler/openai_chat_completions.go b/backend/internal/handler/openai_chat_completions.go index b7ce4c63b..8309729ab 100644 --- a/backend/internal/handler/openai_chat_completions.go +++ b/backend/internal/handler/openai_chat_completions.go @@ -5,6 +5,7 @@ import ( "errors" "net/http" "strconv" + "strings" "time" pkghttputil "github.com/Wei-Shaw/sub2api/internal/pkg/httputil" @@ -89,6 +90,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) { service.SetOpsLatencyMs(c, service.OpsAuthLatencyMsKey, time.Since(requestStart).Milliseconds()) routingStart := time.Now() + c.Request = c.Request.WithContext(service.WithOpenAIFirstOutputStart(c.Request.Context(), routingStart)) userReleaseFunc, acquired := h.acquireResponsesUserSlot(c, subject.UserID, subject.Concurrency, reqStream, &streamStarted, reqLog) if !acquired { @@ -155,6 +157,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) { zap.Int64p("group_id", currentAPIKey.GroupID), ) selectionModel := resolveOpenAIAccountSelectionModel(reqModel, channelMapping) + dispatchModel := selectionModel selectionCtx := openAIAccountShareModeRequestContext(c, currentAPIKey) selectionCtx = openAICompatibleRequestContext(selectionCtx, currentAPIKey) if decision := h.checkCyberPreflightWithContext(selectionCtx, c, reqLog, currentAPIKey, subject, service.ContentModerationProtocolOpenAIChat, reqModel, body); decision != nil && decision.Blocked { @@ -213,6 +216,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) { errorRoutingModel = defaultModel if err == nil && selection != nil { c.Set("openai_chat_completions_fallback_model", defaultModel) + dispatchModel = defaultModel } } if err != nil { @@ -269,10 +273,14 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) { return } - accountReleaseFunc, acquired := h.acquireResponsesAccountSlot(c, currentAPIKey.GroupID, sessionHash, selection, reqStream, &streamStarted, reqLog) + freshAccount, accountReleaseFunc, acquired := h.acquireResponsesAccountSlot(c, selectionCtx, currentAPIKey.GroupID, sessionHash, service.OpenAIAccountDispatchRequirements{ + RequestedModel: dispatchModel, + RequiredTransport: service.OpenAIUpstreamTransportAny, + }, selection, reqStream, &streamStarted, reqLog) if !acquired { return } + account = freshAccount service.SetOpsLatencyMs(c, service.OpsRoutingLatencyMsKey, time.Since(routingStart).Milliseconds()) forwardStart := time.Now() @@ -367,10 +375,14 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) { ) return } + scheduleModel := account.GetMappedModel(dispatchModel) + if result != nil && strings.TrimSpace(result.UpstreamModel) != "" { + scheduleModel = result.UpstreamModel + } if result != nil { - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, result.FirstTokenMs) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, result.FirstTokenMs, scheduleModel) } else { - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, nil) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, nil, scheduleModel) } routeCursor.recordSuccess(apiKey.ID) diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index 478caad4f..d5eec2a4f 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -28,16 +28,21 @@ import ( // OpenAIGatewayHandler handles OpenAI API gateway requests type OpenAIGatewayHandler struct { - gatewayService *service.OpenAIGatewayService - billingCacheService *service.BillingCacheService - apiKeyService *service.APIKeyService - usageRecordWorkerPool *service.UsageRecordWorkerPool - errorPassthroughService *service.ErrorPassthroughService - contentModerationService *service.ContentModerationService - userModerationService *service.UserContentModerationService - concurrencyHelper *ConcurrencyHelper - maxAccountSwitches int - cfg *config.Config + gatewayService *service.OpenAIGatewayService + billingCacheService *service.BillingCacheService + apiKeyService *service.APIKeyService + usageRecordWorkerPool *service.UsageRecordWorkerPool + errorPassthroughService *service.ErrorPassthroughService + contentModerationService *service.ContentModerationService + userModerationService *service.UserContentModerationService + grokMediaEligibilityProber grokMediaEligibilityProber + concurrencyHelper *ConcurrencyHelper + maxAccountSwitches int + cfg *config.Config +} + +type grokMediaEligibilityProber interface { + ProbeMediaEligibility(ctx context.Context, accountID int64) (bool, string, error) } const maxOpenAIFirstOutputTimeoutSwitches = 1 @@ -180,6 +185,11 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { return } reqModel := analysis.Model + imageIntent := service.IsExplicitImageGenerationIntent("/v1/responses", reqModel, body) + if imageIntent && !service.GroupAllowsImageGeneration(apiKey.Group) { + h.errorResponse(c, http.StatusForbidden, "permission_error", service.ImageGenerationPermissionMessage()) + return + } reqStream := analysis.Stream reqLog = reqLog.With(zap.String("model", reqModel), zap.Bool("stream", reqStream)) @@ -224,6 +234,11 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { service.SetOpsLatencyMs(c, service.OpsAuthLatencyMsKey, time.Since(requestStart).Milliseconds()) routingStart := time.Now() + routingCtx := service.WithOpenAIFirstOutputStart(c.Request.Context(), routingStart) + if reqStream { + routingCtx = service.WithOpenAIFirstOutputBudget(routingCtx, h.gatewayService.OpenAIFirstOutputRoutingBudget(body, reqModel)) + } + c.Request = c.Request.WithContext(routingCtx) userReleaseFunc, acquired := h.acquireResponsesUserSlot(c, subject.UserID, subject.Concurrency, reqStream, &streamStarted, reqLog) if !acquired { @@ -261,6 +276,9 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { var lastFailoverErr *service.UpstreamFailoverError for { + if reqStream && h.abortIfOpenAIFirstOutputBudgetExpired(c, streamStarted) { + return + } if !openAIRequestAllowsFailoverReplay(c) { return } @@ -274,8 +292,16 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { if h.rejectIfCyberSessionBlocked(c, currentAPIKey, sessionHashBody, reqModel, cyberBlockFormatResponses) { return } - currentSubscription, subErr := h.gatewayService.ResolveRouteSubscription(c.Request.Context(), currentAPIKey, subscription) + selectionCtx := openAIAccountShareModeRequestContext(c, currentAPIKey) + selectionCtx = openAICompatibleRequestContext(selectionCtx, currentAPIKey) + selectionCtx, cancelSelectionRouting := service.WithOpenAIFirstOutputRoutingDeadline(selectionCtx) + currentSubscription, subErr := h.gatewayService.ResolveRouteSubscription(selectionCtx, currentAPIKey, subscription) + if reqStream && h.abortIfOpenAIFirstOutputBudgetExpired(c, streamStarted) { + cancelSelectionRouting() + return + } if subErr != nil { + cancelSelectionRouting() status, code, message, retryAfter := billingErrorDetails(subErr) if retryAfter > 0 { c.Header("Retry-After", strconv.Itoa(retryAfter)) @@ -283,8 +309,17 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { h.handleStreamingAwareError(c, status, code, message, streamStarted) return } - channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(c.Request.Context(), currentAPIKey.GroupID, reqModel) - if err := h.billingCacheService.CheckBillingEligibility(c.Request.Context(), currentAPIKey.User, currentAPIKey, currentAPIKey.Group, currentSubscription); err != nil { + channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(selectionCtx, currentAPIKey.GroupID, reqModel) + if reqStream && h.abortIfOpenAIFirstOutputBudgetExpired(c, streamStarted) { + cancelSelectionRouting() + return + } + if err := h.billingCacheService.CheckBillingEligibility(selectionCtx, currentAPIKey.User, currentAPIKey, currentAPIKey.Group, currentSubscription); err != nil { + if reqStream && h.abortIfOpenAIFirstOutputBudgetExpired(c, streamStarted) { + cancelSelectionRouting() + return + } + cancelSelectionRouting() reqLog.Info("openai.billing_eligibility_check_failed", zap.Error(err), zap.Int64p("group_id", currentAPIKey.GroupID), @@ -303,20 +338,24 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { zap.Int64p("group_id", currentAPIKey.GroupID), ) selectionModel := resolveOpenAIAccountSelectionModel(reqModel, channelMapping) - selectionCtx := openAIAccountShareModeRequestContext(c, currentAPIKey) - selectionCtx = openAICompatibleRequestContext(selectionCtx, currentAPIKey) routeKey := moderationRouteKey{apiKeyID: currentAPIKey.ID, groupID: apiKeyGroupIDValue(currentAPIKey)} if _, checked := moderatedRoutes[routeKey]; !checked { moderatedRoutes[routeKey] = struct{}{} if decision := h.checkCyberPreflightWithSource(selectionCtx, c, reqLog, currentAPIKey, subject, service.ContentModerationProtocolOpenAIResponses, reqModel, body, analysis); decision != nil && decision.Blocked { + cancelSelectionRouting() h.handleStreamingAwareError(c, contentModerationStatus(decision), cyberPreflightErrorCode(decision), decision.Message, streamStarted) return } if decision := h.checkContentModerationWithSource(selectionCtx, c, reqLog, currentAPIKey, subject, service.ContentModerationProtocolOpenAIResponses, reqModel, body, analysis); decision != nil && decision.Blocked { + cancelSelectionRouting() h.handleStreamingAwareError(c, contentModerationStatus(decision), contentModerationErrorCode(decision), decision.Message, streamStarted) return } } + if reqStream && h.abortIfOpenAIFirstOutputBudgetExpired(c, streamStarted) { + cancelSelectionRouting() + return + } selection, scheduleDecision, err := h.gatewayService.SelectAccountWithCleanRelayScheduler( selectionCtx, c, @@ -331,6 +370,10 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { sessionHashBody, ) if err != nil { + cancelSelectionRouting() + if reqStream && h.abortIfOpenAIFirstOutputBudgetExpired(c, streamStarted) { + return + } if failoverClientGone(c) { reqLog.Info("openai.account_select_aborted_client_disconnected", zap.Error(err)) return @@ -375,6 +418,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { return } if selection == nil || selection.Account == nil { + cancelSelectionRouting() cls := classifyNoAccountErrorFromGin(c, h.gatewayService, currentAPIKey, selectionModel, reqModel, routingPlatform) h.handleStreamingAwareError(c, cls.Status, cls.ErrType, cls.Message, streamStarted) return @@ -392,6 +436,11 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { zap.Float64("load_skew", scheduleDecision.LoadSkew), ) account := selection.Account + if reqStream { + accountBudget := h.gatewayService.OpenAIFirstOutputBudgetForAccount(account, body, reqModel, selectionModel) + c.Request = c.Request.WithContext(service.WithOpenAIFirstOutputBudget(c.Request.Context(), accountBudget)) + selectionCtx = service.WithOpenAIFirstOutputBudget(selectionCtx, accountBudget) + } sessionHash = ensureOpenAIPoolModeSessionHash(sessionHash, account) reqLog.Debug("openai.account_selected", zap.Int64("account_id", account.ID), @@ -402,6 +451,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { if _, checked := moderatedAccounts[account.ID]; !checked { moderatedAccounts[account.ID] = struct{}{} if decision := h.checkUserContentModerationWithSource(selectionCtx, c, reqLog, currentAPIKey, subject, account, service.ContentModerationProtocolOpenAIResponses, reqModel, body, analysis); decision != nil && decision.Blocked { + cancelSelectionRouting() if selection.Acquired && selection.ReleaseFunc != nil { selection.ReleaseFunc() } @@ -409,11 +459,24 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { return } } + dispatchCtx := openAIResponsesDispatchContext(c, selectionCtx, currentAPIKey) + cancelSelectionRouting() + if reqStream && h.abortIfOpenAIFirstOutputBudgetExpired(c, streamStarted) { + if selection.Acquired && selection.ReleaseFunc != nil { + selection.ReleaseFunc() + } + return + } - accountReleaseFunc, acquired := h.acquireResponsesAccountSlot(c, currentAPIKey.GroupID, sessionHash, selection, reqStream, &streamStarted, reqLog) + freshAccount, accountReleaseFunc, acquired := h.acquireResponsesAccountSlot(c, dispatchCtx, currentAPIKey.GroupID, sessionHash, service.OpenAIAccountDispatchRequirements{ + RequestedModel: selectionModel, + RequiredTransport: service.OpenAIUpstreamTransportAny, + RequireCompact: requireCompact, + }, selection, reqStream, &streamStarted, reqLog) if !acquired { return } + account = freshAccount // Forward request service.SetOpsLatencyMs(c, service.OpsRoutingLatencyMsKey, time.Since(routingStart).Milliseconds()) @@ -428,9 +491,9 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { if channelMapping.Mapped { forwardAnalysis = analysis.WithBodyAndModel(forwardBody, channelMapping.MappedModel) } - result, err := h.gatewayService.ForwardWithAnalysis(selectionCtx, c, account, forwardBody, forwardAnalysis) + result, err := h.gatewayService.ForwardWithAnalysis(dispatchCtx, c, account, forwardBody, forwardAnalysis) if service.GetOpsCyberPolicy(c) != nil { - h.gatewayService.MarkCyberSessionBlocked(selectionCtx, service.CyberSessionBlockKey(currentAPIKey.ID, c, sessionHashBody)) + h.gatewayService.MarkCyberSessionBlocked(dispatchCtx, service.CyberSessionBlockKey(currentAPIKey.ID, c, sessionHashBody)) } forwardDurationMs := time.Since(forwardStart).Milliseconds() if accountReleaseFunc != nil { @@ -542,9 +605,9 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { if account.Type == service.AccountTypeOAuth { h.gatewayService.UpdateCodexUsageSnapshotFromHeaders(c.Request.Context(), account.ID, result.ResponseHeaders) } - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, result.FirstTokenMs) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, result.FirstTokenMs, account.GetMappedModel(selectionModel)) } else { - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, nil) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, nil, account.GetMappedModel(selectionModel)) } routeCursor.recordSuccess(apiKey.ID) @@ -555,7 +618,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { // 使用量记录通过有界 worker 池提交，避免请求热路径创建无界 goroutine。 h.submitUsageRecordTask(func(ctx context.Context) { - usageCtx := service.WithAccountShareModeRequestFromContext(ctx, selectionCtx) + usageCtx := service.WithAccountShareModeRequestFromContext(ctx, dispatchCtx) if err := h.gatewayService.RecordUsage(usageCtx, &service.OpenAIRecordUsageInput{ Result: result, APIKey: currentAPIKey, @@ -809,6 +872,7 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) { service.SetOpsLatencyMs(c, service.OpsAuthLatencyMsKey, time.Since(requestStart).Milliseconds()) routingStart := time.Now() + c.Request = c.Request.WithContext(service.WithOpenAIFirstOutputStart(c.Request.Context(), routingStart)) userReleaseFunc, acquired := h.acquireResponsesUserSlot(c, subject.UserID, subject.Concurrency, reqStream, &streamStarted, reqLog) if !acquired { @@ -993,10 +1057,14 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) { return } - accountReleaseFunc, acquired := h.acquireResponsesAccountSlot(c, currentAPIKey.GroupID, sessionHash, selection, reqStream, &streamStarted, reqLog) + freshAccount, accountReleaseFunc, acquired := h.acquireResponsesAccountSlot(c, selectionCtx, currentAPIKey.GroupID, sessionHash, service.OpenAIAccountDispatchRequirements{ + RequestedModel: currentRoutingModel, + RequiredTransport: service.OpenAIUpstreamTransportAny, + }, selection, reqStream, &streamStarted, reqLog) if !acquired { return } + account = freshAccount service.SetOpsLatencyMs(c, service.OpsRoutingLatencyMsKey, time.Since(routingStart).Milliseconds()) forwardStart := time.Now() @@ -1099,9 +1167,9 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) { return } if result != nil { - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, result.FirstTokenMs) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, result.FirstTokenMs, account.GetMappedModel(currentRoutingModel)) } else { - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, nil) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, nil, account.GetMappedModel(currentRoutingModel)) } routeCursor.recordSuccess(apiKey.ID) @@ -1289,26 +1357,57 @@ func (h *OpenAIGatewayHandler) acquireResponsesUserSlot( func (h *OpenAIGatewayHandler) acquireResponsesAccountSlot( c *gin.Context, + selectionCtx context.Context, groupID *int64, sessionHash string, + fallbackRequirements service.OpenAIAccountDispatchRequirements, selection *service.AccountSelectionResult, reqStream bool, streamStarted *bool, reqLog *zap.Logger, -) (func(), bool) { +) (*service.Account, func(), bool) { if selection == nil || selection.Account == nil { h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "api_error", "No available accounts", *streamStarted) - return nil, false + return nil, nil, false } ctx := c.Request.Context() account := selection.Account + dispatchRequirements := fallbackRequirements + if selection.OpenAIDispatchRequirements != nil { + dispatchRequirements = *selection.OpenAIDispatchRequirements + } + dispatchCtx := service.WithAccountShareModeRequestFromContext(ctx, selectionCtx) + finishAcquired := func(release func()) (*service.Account, func(), bool) { + latest, err := h.gatewayService.RevalidateSelectedOpenAIAccountForDispatch( + dispatchCtx, + groupID, + account, + dispatchRequirements, + ) + if err != nil { + if release != nil { + release() + } + reqLog.Info("openai.account_selection_invalidated_before_dispatch", + zap.Int64("account_id", account.ID), + zap.Error(err), + ) + h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "api_error", "Selected account is no longer available, please retry", *streamStarted) + return nil, nil, false + } + selection.Account = latest + if err := h.gatewayService.BindStickySession(ctx, groupID, sessionHash, latest.ID); err != nil { + reqLog.Warn("openai.bind_sticky_session_failed", zap.Int64("account_id", latest.ID), zap.Error(err)) + } + return latest, wrapReleaseOnDone(ctx, release), true + } if selection.Acquired { - return wrapReleaseOnDone(ctx, selection.ReleaseFunc), true + return finishAcquired(selection.ReleaseFunc) } if selection.WaitPlan == nil { h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "api_error", "No available accounts", *streamStarted) - return nil, false + return nil, nil, false } fastReleaseFunc, fastAcquired, err := h.concurrencyHelper.TryAcquireAccountSlot( @@ -1319,13 +1418,10 @@ func (h *OpenAIGatewayHandler) acquireResponsesAccountSlot( if err != nil { reqLog.Warn("openai.account_slot_quick_acquire_failed", zap.Int64("account_id", account.ID), zap.Error(err)) h.handleConcurrencyError(c, err, "account", *streamStarted) - return nil, false + return nil, nil, false } if fastAcquired { - if err := h.gatewayService.BindStickySession(ctx, groupID, sessionHash, account.ID); err != nil { - reqLog.Warn("openai.bind_sticky_session_failed", zap.Int64("account_id", account.ID), zap.Error(err)) - } - return wrapReleaseOnDone(ctx, fastReleaseFunc), true + return finishAcquired(fastReleaseFunc) } canWait, waitErr := h.concurrencyHelper.IncrementAccountWaitCount(ctx, account.ID, selection.WaitPlan.MaxWaiting) @@ -1337,7 +1433,7 @@ func (h *OpenAIGatewayHandler) acquireResponsesAccountSlot( zap.Int("max_waiting", selection.WaitPlan.MaxWaiting), ) h.handleStreamingAwareError(c, http.StatusTooManyRequests, "rate_limit_error", "Too many pending requests, please retry later", *streamStarted) - return nil, false + return nil, nil, false } accountWaitCounted := waitErr == nil && canWait @@ -1360,15 +1456,12 @@ func (h *OpenAIGatewayHandler) acquireResponsesAccountSlot( if err != nil { reqLog.Warn("openai.account_slot_acquire_failed", zap.Int64("account_id", account.ID), zap.Error(err)) h.handleConcurrencyError(c, err, "account", *streamStarted) - return nil, false + return nil, nil, false } // Slot acquired: no longer waiting in queue. releaseWait() - if err := h.gatewayService.BindStickySession(ctx, groupID, sessionHash, account.ID); err != nil { - reqLog.Warn("openai.bind_sticky_session_failed", zap.Int64("account_id", account.ID), zap.Error(err)) - } - return wrapReleaseOnDone(ctx, accountReleaseFunc), true + return finishAcquired(accountReleaseFunc) } // ResponsesWebSocket handles OpenAI Responses API WebSocket ingress endpoint @@ -1405,6 +1498,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, @@ -1425,35 +1540,17 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { _ = wsConn.CloseNow() }() wsConn.SetReadLimit(16 * 1024 * 1024) - - 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 := 30 * time.Second if h.cfg != nil && h.cfg.Gateway.OpenAIWS.ClientFirstMessageTimeoutSeconds > 0 { firstMessageTimeout = time.Duration(h.cfg.Gateway.OpenAIWS.ClientFirstMessageTimeoutSeconds) * time.Second } - 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)) @@ -1561,6 +1658,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { var selection *service.AccountSelectionResult var scheduleDecision service.OpenAIAccountScheduleDecision var selectedAccountShareCtx context.Context + var selectedRoutingModel string var cyberBlockKeyWS string for { if failoverClientGone(c) { @@ -1623,6 +1721,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { ) if selectErr == nil && selection != nil && selection.Account != nil { selectedAccountShareCtx = selectionCtx + selectedRoutingModel = selectionModel break } if failoverClientGone(c) { @@ -1688,6 +1787,27 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { accountReleaseFunc = fastReleaseFunc currentAccountRelease = wrapReleaseOnDone(ctx, accountReleaseFunc) } + dispatchCtx := service.WithAccountShareModeRequestFromContext(ctx, selectedAccountShareCtx) + dispatchRequirements := service.OpenAIAccountDispatchRequirements{ + RequestedModel: selectedRoutingModel, + RequiredTransport: service.OpenAIUpstreamTransportResponsesWebsocketV2, + } + if selection.OpenAIDispatchRequirements != nil { + dispatchRequirements = *selection.OpenAIDispatchRequirements + } + latestAccount, err := h.gatewayService.RevalidateSelectedOpenAIAccountForDispatch( + dispatchCtx, + currentAPIKey.GroupID, + account, + dispatchRequirements, + ) + if err != nil { + reqLog.Info("openai.websocket_selection_invalidated_before_dispatch", zap.Int64("account_id", account.ID), zap.Error(err)) + closeOpenAIClientWS(wsConn, coderws.StatusTryAgainLater, "selected account is no longer available; please reconnect") + return + } + account = latestAccount + accountMaxConcurrency = latestAccount.Concurrency if err := h.gatewayService.BindStickySession(ctx, currentAPIKey.GroupID, sessionHash, account.ID); err != nil { reqLog.Warn("openai.websocket_bind_sticky_session_failed", zap.Int64("account_id", account.ID), zap.Error(err)) } @@ -1718,11 +1838,24 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { if cyberBlockedThisConn { return service.NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, cyberSessionBlockedClientMsg, nil) } + if turn != 1 { + // 防御式清理：避免异常路径下旧槽位覆盖导致泄漏。 + releaseTurnSlots() + } + latest, revalidateErr := h.gatewayService.RevalidateSelectedOpenAIAccountForDispatch( + dispatchCtx, + currentAPIKey.GroupID, + account, + dispatchRequirements, + ) + if revalidateErr != nil { + return service.NewOpenAIWSClientCloseError(coderws.StatusTryAgainLater, "selected account is no longer available; please reconnect", revalidateErr) + } + account = latest + accountMaxConcurrency = latest.Concurrency if turn == 1 { return nil } - // 防御式清理：避免异常路径下旧槽位覆盖导致泄漏。 - releaseTurnSlots() // 非首轮 turn 需要重新抢占并发槽位，避免长连接空闲占槽。 userReleaseFunc, userAcquired, err := h.concurrencyHelper.TryAcquireUserSlotForAPIKey(ctx, subject.UserID, subject.Concurrency, currentAPIKey.ID) if err != nil { @@ -1746,6 +1879,18 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { } currentUserRelease = wrapReleaseOnDone(ctx, userReleaseFunc) currentAccountRelease = wrapReleaseOnDone(ctx, accountReleaseFunc) + latest, revalidateErr = h.gatewayService.RevalidateSelectedOpenAIAccountForDispatch( + dispatchCtx, + currentAPIKey.GroupID, + account, + dispatchRequirements, + ) + if revalidateErr != nil { + releaseTurnSlots() + return service.NewOpenAIWSClientCloseError(coderws.StatusTryAgainLater, "selected account is no longer available; please reconnect", revalidateErr) + } + account = latest + accountMaxConcurrency = latest.Concurrency return nil }, AfterTurn: func(turn int, result *service.OpenAIForwardResult, turnErr error) { @@ -1760,7 +1905,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { if account.Type == service.AccountTypeOAuth { h.gatewayService.UpdateCodexUsageSnapshotFromHeaders(ctx, account.ID, result.ResponseHeaders) } - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, result.FirstTokenMs) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, result.FirstTokenMs, result.UpstreamModel) h.submitUsageRecordTask(func(taskCtx context.Context) { usageCtx := service.WithAccountShareModeRequestFromContext(taskCtx, selectedAccountShareCtx) if err := h.gatewayService.RecordUsage(usageCtx, &service.OpenAIRecordUsageInput{ @@ -1963,8 +2108,15 @@ func (h *OpenAIGatewayHandler) submitUsageRecordTask(task service.UsageRecordTas runUsageRecordTaskSync(task, "handler.openai_gateway.responses", "openai.usage_record_task_panic_recovered") } -// handleConcurrencyError handles concurrency-related errors with proper 429 response +// handleConcurrencyError distinguishes a gateway first-output budget from a +// real concurrency rejection so users and metrics do not misclassify a 504 as +// an account/user rate limit. func (h *OpenAIGatewayHandler) handleConcurrencyError(c *gin.Context, err error, slotType string, streamStarted bool) { + if errors.Is(err, service.ErrOpenAIFirstOutputRoutingBudgetExceeded) { + h.handleStreamingAwareError(c, http.StatusGatewayTimeout, "routing_budget_exhausted", + "Gateway routing budget expired before an upstream attempt could start", streamStarted) + return + } var waitErr *WaitQueueFullError if errors.As(err, &waitErr) { h.handleStreamingAwareError(c, http.StatusTooManyRequests, "rate_limit_error", waitErr.Error(), streamStarted) @@ -1974,11 +2126,42 @@ func (h *OpenAIGatewayHandler) handleConcurrencyError(c *gin.Context, err error, fmt.Sprintf("Concurrency limit exceeded for %s, please retry later", slotType), streamStarted) } +func (h *OpenAIGatewayHandler) abortIfOpenAIFirstOutputBudgetExpired(c *gin.Context, streamStarted bool) bool { + if c == nil || c.Request == nil { + return false + } + remaining, enabled := service.OpenAIFirstOutputBudgetRemaining(c.Request.Context()) + if !enabled || remaining > 0 { + return false + } + h.handleStreamingAwareError(c, http.StatusGatewayTimeout, "routing_budget_exhausted", + "Gateway routing budget expired before an upstream attempt could start", streamStarted) + return true +} + func (h *OpenAIGatewayHandler) handleFailoverExhausted(c *gin.Context, failoverErr *service.UpstreamFailoverError, streamStarted bool) { if failoverErr == nil { h.handleFailoverExhaustedSimple(c, http.StatusBadGateway, streamStarted) return } + if failoverErr.Reason == service.GatewayFailureReasonRoutingBudgetExhausted { + // This is a request-scoped local deadline, not an upstream 504. Preserve + // the same contract used when the budget expires in a concurrency wait. + h.handleStreamingAwareError(c, http.StatusGatewayTimeout, "routing_budget_exhausted", + "Gateway routing budget expired before an upstream attempt could start", streamStarted) + return + } + if failoverErr.IsOpenAIRequestBodyTooLarge() { + service.SetOpsUpstreamError(c, http.StatusRequestEntityTooLarge, service.OpenAIRequestBodyTooLargeClientMessage, "") + h.handleStreamingAwareError( + c, + http.StatusRequestEntityTooLarge, + "invalid_request_error", + service.OpenAIRequestBodyTooLargeClientMessage, + streamStarted, + ) + return + } if failoverErr.IsCredentialFailure() { status := failoverErr.ClientStatusCode if status <= 0 { diff --git a/backend/internal/handler/openai_gateway_handler_test.go b/backend/internal/handler/openai_gateway_handler_test.go index e35c21a6d..8d0e7dcd0 100644 --- a/backend/internal/handler/openai_gateway_handler_test.go +++ b/backend/internal/handler/openai_gateway_handler_test.go @@ -10,6 +10,7 @@ import ( "testing" "time" + "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" pkghttputil "github.com/Wei-Shaw/sub2api/internal/pkg/httputil" "github.com/Wei-Shaw/sub2api/internal/server/middleware" "github.com/Wei-Shaw/sub2api/internal/service" @@ -190,6 +191,27 @@ func TestOpenAIImagesRequestFailureReturnsAccurateClientStatus(t *testing.T) { require.Equal(t, "unsupported image option", gjson.Get(recorder.Body.String(), "error.message").String()) } +func TestOpenAIRequestBodyTooLargeFailoverExhaustedReturnsSanitized413(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + h := &OpenAIGatewayHandler{} + + h.handleFailoverExhausted(c, &service.UpstreamFailoverError{ + StatusCode: http.StatusRequestEntityTooLarge, + ResponseBody: []byte(`{"error":{"message":"proxy internal.example leaked tenant-secret"}}`), + Reason: service.GatewayFailureReason("openai_request_body_too_large"), + ClientStatusCode: http.StatusRequestEntityTooLarge, + ClientMessage: service.OpenAIRequestBodyTooLargeClientMessage, + }, false) + + require.Equal(t, http.StatusRequestEntityTooLarge, recorder.Code) + require.Equal(t, "invalid_request_error", gjson.Get(recorder.Body.String(), "error.type").String()) + require.Equal(t, service.OpenAIRequestBodyTooLargeClientMessage, gjson.Get(recorder.Body.String(), "error.message").String()) + require.NotContains(t, recorder.Body.String(), "internal.example") + require.NotContains(t, recorder.Body.String(), "tenant-secret") +} + func TestAppendOpenAIProxyLogFields(t *testing.T) { base := []zap.Field{zap.Int64("account_id", 7)} @@ -667,6 +689,30 @@ func TestOpenAIResponses_RejectsHTTPContinuationPreviousResponseID(t *testing.T) require.Contains(t, w.Body.String(), "previous_response_id") } +func TestOpenAIResponses_RejectsExplicitImageIntentWhenGroupDisallowsImages(t *testing.T) { + gin.SetMode(gin.TestMode) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", strings.NewReader( + `{"model":"gpt-5.4","input":"generate an image of a lighthouse","tools":[{"type":"image_generation"}]}`, + )) + c.Request.Header.Set("Content-Type", "application/json") + groupID := int64(2) + c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{ + ID: 101, + GroupID: &groupID, + Group: &service.Group{ID: groupID, AllowImageGeneration: false}, + User: &service.User{ID: 1}, + }) + c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: 1, Concurrency: 1}) + + newOpenAIHandlerForPreviousResponseIDValidation(t, nil).Responses(c) + + require.Equal(t, http.StatusForbidden, w.Code) + require.Contains(t, w.Body.String(), service.ImageGenerationPermissionMessage()) +} + func TestOpenAIResponses_FunctionCallOutputHTTPGuidanceDoesNotSuggestPreviousResponseReuse(t *testing.T) { gin.SetMode(gin.TestMode) @@ -1036,3 +1082,38 @@ func TestOpenAIFirstOutputFailoverExhaustedAllowsOnlyOneAccountSwitch(t *testing require.False(t, openAIFirstOutputFailoverExhausted(&service.UpstreamFailoverError{}, &switchCount)) require.False(t, openAIFirstOutputFailoverExhausted(nil, &switchCount)) } + +func TestOpenAIResponsesDispatchContextDetachesRoutingCancellation(t *testing.T) { + gin.SetMode(gin.TestMode) + requestCtx, cancelRequest := context.WithCancel(context.Background()) + t.Cleanup(cancelRequest) + requestCtx = service.WithOpenAIFirstOutputStart(requestCtx, time.Now()) + requestCtx = service.WithOpenAIFirstOutputBudget(requestCtx, time.Minute) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil).WithContext(requestCtx) + apiKey := &service.APIKey{ + ID: 22, + UserID: 11, + Group: &service.Group{Platform: service.PlatformGrok}, + } + routingCtx := openAIAccountShareModeRequestContext(c, apiKey) + routingCtx = openAICompatibleRequestContext(routingCtx, apiKey) + routingCtx, cancelRouting := service.WithOpenAIFirstOutputRoutingDeadline(routingCtx) + dispatchCtx := openAIResponsesDispatchContext(c, routingCtx, apiKey) + + cancelRouting() + require.ErrorIs(t, routingCtx.Err(), context.Canceled) + require.NoError(t, dispatchCtx.Err()) + shareMode, ok := service.AccountShareModeRequestFromContext(dispatchCtx) + require.True(t, ok) + require.Equal(t, int64(11), shareMode.UserID) + require.Equal(t, int64(22), shareMode.APIKeyID) + require.Equal(t, service.PlatformGrok, dispatchCtx.Value(ctxkey.ForcePlatform)) + _, budgetEnabled := service.OpenAIFirstOutputBudgetRemaining(dispatchCtx) + require.True(t, budgetEnabled) + + cancelRequest() + require.ErrorIs(t, dispatchCtx.Err(), context.Canceled) +} diff --git a/backend/internal/handler/openai_images.go b/backend/internal/handler/openai_images.go index 8436f6e00..0f8104765 100644 --- a/backend/internal/handler/openai_images.go +++ b/backend/internal/handler/openai_images.go @@ -235,10 +235,15 @@ routeLoop: return } - accountReleaseFunc, acquired := h.acquireResponsesAccountSlot(c, currentAPIKey.GroupID, sessionHash, selection, parsed.Stream, &streamStarted, reqLog) + freshAccount, accountReleaseFunc, acquired := h.acquireResponsesAccountSlot(c, selectionCtx, currentAPIKey.GroupID, sessionHash, service.OpenAIAccountDispatchRequirements{ + RequestedModel: selectionModel, + RequiredTransport: service.OpenAIUpstreamTransportHTTPSSE, + RequiredImageCapability: parsed.RequiredCapability, + }, selection, parsed.Stream, &streamStarted, reqLog) if !acquired { return } + account = freshAccount service.SetOpsLatencyMs(c, service.OpsRoutingLatencyMsKey, time.Since(routingStart).Milliseconds()) if !parsed.Stream && !jsonKeepaliveStarted { @@ -343,9 +348,9 @@ routeLoop: if account.Type == service.AccountTypeOAuth { h.gatewayService.UpdateCodexUsageSnapshotFromHeaders(c.Request.Context(), account.ID, result.ResponseHeaders) } - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, result.FirstTokenMs) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, result.FirstTokenMs, account.GetMappedModel(selectionModel)) } else { - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, nil) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, nil, account.GetMappedModel(selectionModel)) } routeCursor.recordSuccess(apiKey.ID) diff --git a/backend/internal/handler/ops_error_logger.go b/backend/internal/handler/ops_error_logger.go index f7a7bab75..25ad56b9e 100644 --- a/backend/internal/handler/ops_error_logger.go +++ b/backend/internal/handler/ops_error_logger.go @@ -26,10 +26,11 @@ import ( ) const ( - opsModelKey = "ops_model" - opsStreamKey = "ops_stream" - opsRequestBodyKey = "ops_request_body" - opsAccountIDKey = "ops_account_id" + opsModelKey = "ops_model" + opsStreamKey = "ops_stream" + opsRequestBodyKey = "ops_request_body" + opsAccountIDKey = "ops_account_id" + opsRoutingCapacityLimitedKey = "ops_routing_capacity_limited" opsUpstreamModelKey = "ops_upstream_model" opsRequestTypeKey = "ops_request_type" @@ -431,6 +432,25 @@ func setOpsSelectedAccount(c *gin.Context, accountID int64, platform ...string) } } +func markOpsRoutingCapacityLimited(c *gin.Context) { + if c == nil { + return + } + c.Set(opsRoutingCapacityLimitedKey, true) +} + +func isOpsRoutingCapacityLimited(c *gin.Context) bool { + if c == nil { + return false + } + value, exists := c.Get(opsRoutingCapacityLimitedKey) + if !exists { + return false + } + marked, _ := value.(bool) + return marked +} + type opsCaptureWriter struct { gin.ResponseWriter limit int @@ -893,6 +913,10 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc { phase := classifyOpsPhase(normalizedType, parsed.Message, parsed.Code) isBusinessLimited := classifyOpsIsBusinessLimited(normalizedType, phase, parsed.Code, status, parsed.Message) + if isOpsRoutingCapacityLimited(c) { + phase = "routing" + isBusinessLimited = false + } errorOwner := classifyOpsErrorOwner(phase, parsed.Message) errorSource := classifyOpsErrorSource(phase, parsed.Message) diff --git a/backend/internal/handler/setting_handler.go b/backend/internal/handler/setting_handler.go index f31fcfce7..bb5bb969d 100644 --- a/backend/internal/handler/setting_handler.go +++ b/backend/internal/handler/setting_handler.go @@ -85,12 +85,13 @@ func (h *SettingHandler) GetPublicSettings(c *gin.Context) { OpenAIAccountLevels: openAIAccountLevelsToDTO(settings.OpenAIAccountLevels), - AffiliateEnabled: settings.AffiliateEnabled, - RiskControlEnabled: settings.RiskControlEnabled, - InvoiceManagementEnabled: settings.InvoiceManagementEnabled, - WithdrawalManagementEnabled: settings.WithdrawalManagementEnabled, - WithdrawalRateLimitWindowDays: settings.WithdrawalRateLimitWindowDays, - WithdrawalRateLimitMax: settings.WithdrawalRateLimitMax, + AffiliateEnabled: settings.AffiliateEnabled, + RiskControlEnabled: settings.RiskControlEnabled, + InvoiceManagementEnabled: settings.InvoiceManagementEnabled, + WithdrawalManagementEnabled: settings.WithdrawalManagementEnabled, + WithdrawalRateLimitWindowDays: settings.WithdrawalRateLimitWindowDays, + WithdrawalRateLimitMax: settings.WithdrawalRateLimitMax, + WithdrawalRateLimitExemptAmount: settings.WithdrawalRateLimitExemptAmount, }) } diff --git a/backend/internal/handler/usage_handler.go b/backend/internal/handler/usage_handler.go index 4806218b2..5dea96ac3 100644 --- a/backend/internal/handler/usage_handler.go +++ b/backend/internal/handler/usage_handler.go @@ -393,63 +393,149 @@ func (h *UsageHandler) Stats(c *gin.Context) { response.Success(c, stats) } -// parseUserTimeRange parses start_date, end_date query parameters for user dashboard -// Uses user's timezone if provided, otherwise falls back to server timezone -func parseUserTimeRange(c *gin.Context) (time.Time, time.Time) { - userTZ := c.Query("timezone") // Get user's timezone from request - now := timezone.NowInUserLocation(userTZ) +type userDashboardTimeRange struct { + startTime time.Time + endTime time.Time + location *time.Location +} + +func loadUserDashboardLocation(c *gin.Context) (*time.Location, error) { + timezoneName := strings.TrimSpace(c.Query("timezone")) + if timezoneName == "" { + return timezone.Location(), nil + } + + location, err := time.LoadLocation(timezoneName) + if err != nil { + return nil, fmt.Errorf("invalid timezone %q, use an IANA timezone name", timezoneName) + } + return location, nil +} + +func parseUserDashboardTimestamp(raw string, location *time.Location, field string) (time.Time, error) { + raw = strings.TrimSpace(raw) + if timestamp, err := time.Parse(time.RFC3339Nano, raw); err == nil { + return timestamp.In(location), nil + } + for _, layout := range []string{"2006-01-02T15:04:05", "2006-01-02T15:04"} { + if timestamp, err := time.ParseInLocation(layout, raw, location); err == nil { + return timestamp, nil + } + } + return time.Time{}, fmt.Errorf("Invalid %s format, use RFC3339 or YYYY-MM-DDTHH:mm:ss", field) +} + +func parseUserDashboardExactTimeRange(c *gin.Context, location *time.Location) (*time.Time, *time.Time, bool, error) { + startRaw := strings.TrimSpace(c.Query("start_time")) + endRaw := strings.TrimSpace(c.Query("end_time")) + if startRaw == "" && endRaw == "" { + return nil, nil, false, nil + } + if startRaw == "" || endRaw == "" { + return nil, nil, true, fmt.Errorf("start_time and end_time must be provided together") + } + + startTime, err := parseUserDashboardTimestamp(startRaw, location, "start_time") + if err != nil { + return nil, nil, true, err + } + endTime, err := parseUserDashboardTimestamp(endRaw, location, "end_time") + if err != nil { + return nil, nil, true, err + } + if !endTime.After(startTime) { + return nil, nil, true, fmt.Errorf("end_time must be after start_time") + } + if strings.TrimSpace(c.Query("start_date")) != "" || strings.TrimSpace(c.Query("end_date")) != "" { + return nil, nil, true, fmt.Errorf("start_time/end_time cannot be combined with start_date/end_date") + } + return &startTime, &endTime, true, nil +} + +func startOfDayInLocation(value time.Time, location *time.Location) time.Time { + value = value.In(location) + return time.Date(value.Year(), value.Month(), value.Day(), 0, 0, 0, 0, location) +} + +// parseUserTimeRange parses exact timestamps or legacy date parameters for user dashboard. +// Invalid legacy dates retain their historical fallback behavior for compatibility. +func parseUserTimeRange(c *gin.Context) (userDashboardTimeRange, error) { + location, err := loadUserDashboardLocation(c) + if err != nil { + return userDashboardTimeRange{}, err + } + startExact, endExact, hasExactTime, err := parseUserDashboardExactTimeRange(c, location) + if err != nil { + return userDashboardTimeRange{}, err + } + if hasExactTime { + return userDashboardTimeRange{startTime: *startExact, endTime: *endExact, location: location}, nil + } + + now := time.Now().In(location) startDate := c.Query("start_date") endDate := c.Query("end_date") var startTime, endTime time.Time if startDate != "" { - if t, err := timezone.ParseInUserLocation("2006-01-02", startDate, userTZ); err == nil { + if t, err := time.ParseInLocation("2006-01-02", startDate, location); err == nil { startTime = t } else { - startTime = timezone.StartOfDayInUserLocation(now.AddDate(0, 0, -7), userTZ) + startTime = startOfDayInLocation(now.AddDate(0, 0, -7), location) } } else { - startTime = timezone.StartOfDayInUserLocation(now.AddDate(0, 0, -7), userTZ) + startTime = startOfDayInLocation(now.AddDate(0, 0, -7), location) } if endDate != "" { - if t, err := timezone.ParseInUserLocation("2006-01-02", endDate, userTZ); err == nil { - endTime = t.Add(24 * time.Hour) // Include the end date + if t, err := time.ParseInLocation("2006-01-02", endDate, location); err == nil { + endTime = t.AddDate(0, 0, 1) // Include the end date using calendar-day semantics. } else { - endTime = timezone.StartOfDayInUserLocation(now.AddDate(0, 0, 1), userTZ) + endTime = startOfDayInLocation(now.AddDate(0, 0, 1), location) } } else { - endTime = timezone.StartOfDayInUserLocation(now.AddDate(0, 0, 1), userTZ) + endTime = startOfDayInLocation(now.AddDate(0, 0, 1), location) } - return startTime, endTime + return userDashboardTimeRange{startTime: startTime, endTime: endTime, location: location}, nil } -func parseUserDashboardTimeRangeStrict(c *gin.Context) (time.Time, time.Time, error) { - userTZ := c.Query("timezone") - now := timezone.NowInUserLocation(userTZ) - startTime := timezone.StartOfDayInUserLocation(now.AddDate(0, 0, -7), userTZ) - endTime := timezone.StartOfDayInUserLocation(now.AddDate(0, 0, 1), userTZ) +func parseUserDashboardTimeRangeStrict(c *gin.Context) (userDashboardTimeRange, error) { + location, err := loadUserDashboardLocation(c) + if err != nil { + return userDashboardTimeRange{}, err + } + startExact, endExact, hasExactTime, err := parseUserDashboardExactTimeRange(c, location) + if err != nil { + return userDashboardTimeRange{}, err + } + if hasExactTime { + return userDashboardTimeRange{startTime: *startExact, endTime: *endExact, location: location}, nil + } + + now := time.Now().In(location) + startTime := startOfDayInLocation(now.AddDate(0, 0, -7), location) + endTime := startOfDayInLocation(now.AddDate(0, 0, 1), location) if startDate := strings.TrimSpace(c.Query("start_date")); startDate != "" { - t, err := timezone.ParseInUserLocation("2006-01-02", startDate, userTZ) + t, err := time.ParseInLocation("2006-01-02", startDate, location) if err != nil { - return time.Time{}, time.Time{}, fmt.Errorf("invalid start_date format, use YYYY-MM-DD") + return userDashboardTimeRange{}, fmt.Errorf("invalid start_date format, use YYYY-MM-DD") } startTime = t } if endDate := strings.TrimSpace(c.Query("end_date")); endDate != "" { - t, err := timezone.ParseInUserLocation("2006-01-02", endDate, userTZ) + t, err := time.ParseInLocation("2006-01-02", endDate, location) if err != nil { - return time.Time{}, time.Time{}, fmt.Errorf("invalid end_date format, use YYYY-MM-DD") + return userDashboardTimeRange{}, fmt.Errorf("invalid end_date format, use YYYY-MM-DD") } endTime = t.AddDate(0, 0, 1) } if !endTime.After(startTime) { - return time.Time{}, time.Time{}, fmt.Errorf("end_date must be greater than or equal to start_date") + return userDashboardTimeRange{}, fmt.Errorf("end_date must be greater than or equal to start_date") } - return startTime, endTime, nil + return userDashboardTimeRange{startTime: startTime, endTime: endTime, location: location}, nil } // DashboardStats handles getting user dashboard statistics @@ -479,10 +565,21 @@ func (h *UsageHandler) DashboardTrend(c *gin.Context) { return } - startTime, endTime := parseUserTimeRange(c) + timeRange, err := parseUserTimeRange(c) + if err != nil { + response.BadRequest(c, err.Error()) + return + } granularity := c.DefaultQuery("granularity", "day") - trend, err := h.usageService.GetUserUsageTrendByUserID(c.Request.Context(), subject.UserID, startTime, endTime, granularity) + trend, err := h.usageService.GetUserUsageTrendByUserID( + c.Request.Context(), + subject.UserID, + timeRange.startTime, + timeRange.endTime, + granularity, + timeRange.location, + ) if err != nil { response.ErrorFrom(c, err) return @@ -490,8 +587,8 @@ func (h *UsageHandler) DashboardTrend(c *gin.Context) { response.Success(c, gin.H{ "trend": trend, - "start_date": startTime.Format("2006-01-02"), - "end_date": endTime.Add(-24 * time.Hour).Format("2006-01-02"), + "start_date": timeRange.startTime.In(timeRange.location).Format("2006-01-02"), + "end_date": timeRange.endTime.In(timeRange.location).Add(-time.Nanosecond).Format("2006-01-02"), "granularity": granularity, }) } @@ -505,9 +602,13 @@ func (h *UsageHandler) DashboardModels(c *gin.Context) { return } - startTime, endTime := parseUserTimeRange(c) + timeRange, err := parseUserTimeRange(c) + if err != nil { + response.BadRequest(c, err.Error()) + return + } - stats, err := h.usageService.GetUserModelStats(c.Request.Context(), subject.UserID, startTime, endTime) + stats, err := h.usageService.GetUserModelStats(c.Request.Context(), subject.UserID, timeRange.startTime, timeRange.endTime) if err != nil { response.ErrorFrom(c, err) return @@ -515,8 +616,8 @@ func (h *UsageHandler) DashboardModels(c *gin.Context) { response.Success(c, gin.H{ "models": stats, - "start_date": startTime.Format("2006-01-02"), - "end_date": endTime.Add(-24 * time.Hour).Format("2006-01-02"), + "start_date": timeRange.startTime.In(timeRange.location).Format("2006-01-02"), + "end_date": timeRange.endTime.In(timeRange.location).Add(-time.Nanosecond).Format("2006-01-02"), }) } @@ -580,7 +681,7 @@ func (h *UsageHandler) DashboardAccountSharing(c *gin.Context) { return } - startTime, endTime, err := parseUserDashboardTimeRangeStrict(c) + timeRange, err := parseUserDashboardTimeRangeStrict(c) if err != nil { response.BadRequest(c, err.Error()) return @@ -593,11 +694,22 @@ func (h *UsageHandler) DashboardAccountSharing(c *gin.Context) { return } - stats, err := h.usageService.GetUserAccountSharingDashboard(c.Request.Context(), subject.UserID, startTime, endTime, granularity) + stats, err := h.usageService.GetUserAccountSharingDashboard( + c.Request.Context(), + subject.UserID, + timeRange.startTime, + timeRange.endTime, + granularity, + timeRange.location, + ) if err != nil { response.ErrorFrom(c, err) return } + if stats != nil { + stats.StartDate = timeRange.startTime.In(timeRange.location).Format("2006-01-02") + stats.EndDate = timeRange.endTime.In(timeRange.location).Add(-time.Nanosecond).Format("2006-01-02") + } response.Success(c, stats) } diff --git a/backend/internal/handler/usage_handler_dashboard_time_range_test.go b/backend/internal/handler/usage_handler_dashboard_time_range_test.go new file mode 100644 index 000000000..e6115256c --- /dev/null +++ b/backend/internal/handler/usage_handler_dashboard_time_range_test.go @@ -0,0 +1,297 @@ +package handler + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "net/url" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/usagestats" + 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 userDashboardTimeRangeRepoCapture struct { + service.UsageLogRepository + trendStartTime time.Time + trendEndTime time.Time + modelStartTime time.Time + modelEndTime time.Time + accountSharingStartTime time.Time + accountSharingEndTime time.Time + trendLocation *time.Location + accountSharingLocation *time.Location +} + +func (r *userDashboardTimeRangeRepoCapture) GetUserUsageTrendByUserID( + _ context.Context, + _ int64, + startTime, endTime time.Time, + _ string, + location *time.Location, +) ([]usagestats.TrendDataPoint, error) { + r.trendStartTime = startTime + r.trendEndTime = endTime + r.trendLocation = location + return []usagestats.TrendDataPoint{}, nil +} + +func (r *userDashboardTimeRangeRepoCapture) GetUserModelStats( + _ context.Context, + _ int64, + startTime, endTime time.Time, +) ([]usagestats.ModelStat, error) { + r.modelStartTime = startTime + r.modelEndTime = endTime + return []usagestats.ModelStat{}, nil +} + +func (r *userDashboardTimeRangeRepoCapture) GetUserAccountSharingDashboard( + _ context.Context, + _ int64, + startTime, endTime time.Time, + _ string, + location *time.Location, +) (*usagestats.AccountSharingDashboardStats, error) { + r.accountSharingStartTime = startTime + r.accountSharingEndTime = endTime + r.accountSharingLocation = location + return &usagestats.AccountSharingDashboardStats{}, nil +} + +func newUserDashboardTimeRangeTestRouter(repo *userDashboardTimeRangeRepoCapture) *gin.Engine { + gin.SetMode(gin.TestMode) + usageService := service.NewUsageService(repo, nil, nil, nil) + handler := NewUsageHandler(usageService, nil) + router := gin.New() + router.Use(func(c *gin.Context) { + c.Set(string(middleware2.ContextKeyUser), middleware2.AuthSubject{UserID: 42}) + c.Next() + }) + router.GET("/usage/dashboard/trend", handler.DashboardTrend) + router.GET("/usage/dashboard/models", handler.DashboardModels) + router.GET("/usage/dashboard/account-sharing", handler.DashboardAccountSharing) + return router +} + +func TestUserDashboardEndpointsUseExactTimeRange(t *testing.T) { + repo := &userDashboardTimeRangeRepoCapture{} + router := newUserDashboardTimeRangeTestRouter(repo) + wantStart := time.Date(2026, time.July, 22, 12, 34, 56, 789000000, time.UTC) + wantEnd := wantStart.Add(24 * time.Hour) + query := url.Values{ + "start_time": {wantStart.Format(time.RFC3339Nano)}, + "end_time": {wantEnd.Format(time.RFC3339Nano)}, + "timezone": {"Asia/Shanghai"}, + }.Encode() + + tests := []struct { + name string + path string + captured func() (time.Time, time.Time) + }{ + { + name: "trend", + path: "/usage/dashboard/trend", + captured: func() (time.Time, time.Time) { + return repo.trendStartTime, repo.trendEndTime + }, + }, + { + name: "models", + path: "/usage/dashboard/models", + captured: func() (time.Time, time.Time) { + return repo.modelStartTime, repo.modelEndTime + }, + }, + { + name: "account sharing", + path: "/usage/dashboard/account-sharing", + captured: func() (time.Time, time.Time) { + return repo.accountSharingStartTime, repo.accountSharingEndTime + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodGet, tt.path+"?"+query, nil) + router.ServeHTTP(recorder, request) + + require.Equal(t, http.StatusOK, recorder.Code) + gotStart, gotEnd := tt.captured() + require.True(t, gotStart.Equal(wantStart), "start time = %s, want %s", gotStart, wantStart) + require.True(t, gotEnd.Equal(wantEnd), "end time = %s, want %s", gotEnd, wantEnd) + require.Equal(t, "Asia/Shanghai", gotStart.Location().String()) + require.Equal(t, "Asia/Shanghai", gotEnd.Location().String()) + }) + } + require.Equal(t, "Asia/Shanghai", repo.trendLocation.String()) + require.Equal(t, "Asia/Shanghai", repo.accountSharingLocation.String()) +} + +func assertUserDashboardEndpointsReject(t *testing.T, query url.Values, message string) { + t.Helper() + repo := &userDashboardTimeRangeRepoCapture{} + router := newUserDashboardTimeRangeTestRouter(repo) + + for _, path := range []string{ + "/usage/dashboard/trend", + "/usage/dashboard/models", + "/usage/dashboard/account-sharing", + } { + t.Run(path, func(t *testing.T) { + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodGet, path+"?"+query.Encode(), nil) + router.ServeHTTP(recorder, request) + + require.Equal(t, http.StatusBadRequest, recorder.Code) + require.Contains(t, recorder.Body.String(), message) + }) + } +} + +func TestUserDashboardEndpointsRejectIncompleteExactTimeRange(t *testing.T) { + assertUserDashboardEndpointsReject(t, url.Values{ + "start_time": {"2026-07-22T12:34:56Z"}, + "start_date": {"2026-07-01"}, + "end_date": {"2026-07-07"}, + "timezone": {"Asia/Shanghai"}, + }, "start_time and end_time must be provided together") +} + +func TestUserDashboardEndpointsRejectMixedExactAndCalendarRanges(t *testing.T) { + assertUserDashboardEndpointsReject(t, url.Values{ + "start_time": {"2026-07-22T12:34:56Z"}, + "end_time": {"2026-07-23T12:34:56Z"}, + "start_date": {"2026-07-01"}, + "end_date": {"2026-07-07"}, + "timezone": {"Asia/Shanghai"}, + }, "start_time/end_time cannot be combined with start_date/end_date") +} + +func TestUserDashboardEndpointsRejectInvalidTimezone(t *testing.T) { + assertUserDashboardEndpointsReject(t, url.Values{ + "start_time": {"2026-07-22T12:34:56Z"}, + "end_time": {"2026-07-23T12:34:56Z"}, + "timezone": {"Mars/Olympus_Mons"}, + }, "invalid timezone") +} + +func TestUserDashboardEndpointsRejectReversedExactRange(t *testing.T) { + assertUserDashboardEndpointsReject(t, url.Values{ + "start_time": {"2026-07-23T12:34:56Z"}, + "end_time": {"2026-07-22T12:34:56Z"}, + "timezone": {"Asia/Shanghai"}, + }, "end_time must be after start_time") +} + +func TestUserDashboardTrendKeepsCalendarDateRangeCompatibility(t *testing.T) { + repo := &userDashboardTimeRangeRepoCapture{} + router := newUserDashboardTimeRangeTestRouter(repo) + query := url.Values{ + "start_date": {"2026-07-01"}, + "end_date": {"2026-07-07"}, + "timezone": {"Asia/Shanghai"}, + }.Encode() + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodGet, "/usage/dashboard/trend?"+query, nil) + + router.ServeHTTP(recorder, request) + + require.Equal(t, http.StatusOK, recorder.Code) + wantLocation, err := time.LoadLocation("Asia/Shanghai") + require.NoError(t, err) + require.Equal(t, time.Date(2026, time.July, 1, 0, 0, 0, 0, wantLocation), repo.trendStartTime) + require.Equal(t, time.Date(2026, time.July, 8, 0, 0, 0, 0, wantLocation), repo.trendEndTime) + require.Equal(t, wantLocation, repo.trendLocation) +} + +func TestUserDashboardTrendKeepsCalendarDayAcrossDST(t *testing.T) { + repo := &userDashboardTimeRangeRepoCapture{} + router := newUserDashboardTimeRangeTestRouter(repo) + query := url.Values{ + "start_date": {"2026-03-08"}, + "end_date": {"2026-03-08"}, + "timezone": {"America/New_York"}, + }.Encode() + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodGet, "/usage/dashboard/trend?"+query, nil) + + router.ServeHTTP(recorder, request) + + require.Equal(t, http.StatusOK, recorder.Code) + wantLocation, err := time.LoadLocation("America/New_York") + require.NoError(t, err) + require.Equal(t, time.Date(2026, time.March, 8, 0, 0, 0, 0, wantLocation), repo.trendStartTime) + require.Equal(t, time.Date(2026, time.March, 9, 0, 0, 0, 0, wantLocation), repo.trendEndTime) + require.Equal(t, 23*time.Hour, repo.trendEndTime.Sub(repo.trendStartTime)) +} + +func decodeUserDashboardResponseRange(t *testing.T, recorder *httptest.ResponseRecorder) (string, string) { + t.Helper() + var body struct { + Data struct { + StartDate string `json:"start_date"` + EndDate string `json:"end_date"` + } `json:"data"` + } + require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &body)) + return body.Data.StartDate, body.Data.EndDate +} + +func TestUserDashboardEndpointsFormatExactRangeInClientTimezone(t *testing.T) { + repo := &userDashboardTimeRangeRepoCapture{} + router := newUserDashboardTimeRangeTestRouter(repo) + query := url.Values{ + "start_time": {"2026-07-21T16:30:00Z"}, + "end_time": {"2026-07-22T16:30:00Z"}, + "timezone": {"Asia/Shanghai"}, + }.Encode() + + for _, path := range []string{ + "/usage/dashboard/trend", + "/usage/dashboard/models", + "/usage/dashboard/account-sharing", + } { + t.Run(path, func(t *testing.T) { + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodGet, path+"?"+query, nil) + router.ServeHTTP(recorder, request) + + require.Equal(t, http.StatusOK, recorder.Code) + startDate, endDate := decodeUserDashboardResponseRange(t, recorder) + require.Equal(t, "2026-07-22", startDate) + require.Equal(t, "2026-07-23", endDate) + }) + } +} + +func TestUserDashboardExactRangeRemainsTwentyFourHoursAcrossDST(t *testing.T) { + repo := &userDashboardTimeRangeRepoCapture{} + router := newUserDashboardTimeRangeTestRouter(repo) + query := url.Values{ + "start_time": {"2026-03-08T00:30:00-05:00"}, + "end_time": {"2026-03-09T01:30:00-04:00"}, + "timezone": {"America/New_York"}, + }.Encode() + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodGet, "/usage/dashboard/trend?"+query, nil) + + router.ServeHTTP(recorder, request) + + require.Equal(t, http.StatusOK, recorder.Code) + require.Equal(t, 24*time.Hour, repo.trendEndTime.Sub(repo.trendStartTime)) + require.Equal(t, "America/New_York", repo.trendStartTime.Location().String()) + require.Equal(t, 0, repo.trendStartTime.Hour()) + require.Equal(t, 1, repo.trendEndTime.Hour()) + startDate, endDate := decodeUserDashboardResponseRange(t, recorder) + require.Equal(t, "2026-03-08", startDate) + require.Equal(t, "2026-03-09", endDate) +} diff --git a/backend/internal/handler/user_account_agent_identity_import_test.go b/backend/internal/handler/user_account_agent_identity_import_test.go new file mode 100644 index 000000000..c746cc979 --- /dev/null +++ b/backend/internal/handler/user_account_agent_identity_import_test.go @@ -0,0 +1,76 @@ +package handler + +import ( + "testing" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/stretchr/testify/require" +) + +func TestResolveUserOpenAICredentialImportMode(t *testing.T) { + agentIdentity := service.AccountCredentialImportSource{ + Kind: service.AccountCredentialImportKindOpenAIAgentIdentity, + Platform: service.PlatformOpenAI, + } + oauth := service.AccountCredentialImportSource{ + Kind: service.AccountCredentialImportKindOAuthCredentials, + Platform: service.PlatformOpenAI, + } + + tests := []struct { + name string + req importUserAccountCredentialsRequest + sources []service.AccountCredentialImportSource + wantAgent bool + wantErr bool + }{ + { + name: "declared Agent Identity", + req: importUserAccountCredentialsRequest{Platform: service.PlatformOpenAI, OpenAIAuthMode: userOpenAIAuthModeAgentIdentity}, + sources: []service.AccountCredentialImportSource{agentIdentity}, + wantAgent: true, + }, + { + name: "legacy client infers Agent Identity", + req: importUserAccountCredentialsRequest{Platform: service.PlatformOpenAI}, + sources: []service.AccountCredentialImportSource{agentIdentity}, + wantAgent: true, + }, + { + name: "declared Agent Identity rejects OAuth", + req: importUserAccountCredentialsRequest{Platform: service.PlatformOpenAI, OpenAIAuthMode: userOpenAIAuthModeAgentIdentity}, + sources: []service.AccountCredentialImportSource{oauth}, + wantErr: true, + }, + { + name: "declared OAuth rejects Agent Identity", + req: importUserAccountCredentialsRequest{Platform: service.PlatformOpenAI, OpenAIAuthMode: userOpenAIAuthModeOAuth}, + sources: []service.AccountCredentialImportSource{agentIdentity}, + wantErr: true, + }, + { + name: "mixed credentials are rejected", + req: importUserAccountCredentialsRequest{Platform: service.PlatformOpenAI}, + sources: []service.AccountCredentialImportSource{agentIdentity, oauth}, + wantErr: true, + }, + { + name: "wrong platform is rejected", + req: importUserAccountCredentialsRequest{Platform: service.PlatformAnthropic, OpenAIAuthMode: userOpenAIAuthModeAgentIdentity}, + sources: []service.AccountCredentialImportSource{agentIdentity}, + wantErr: true, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + gotAgent, err := resolveUserOpenAICredentialImportMode(test.req, test.sources) + if test.wantErr { + require.Error(t, err) + return + } + require.NoError(t, err) + require.Equal(t, test.wantAgent, gotAgent) + }) + } +} diff --git a/backend/internal/handler/user_account_handler.go b/backend/internal/handler/user_account_handler.go index 33ba640d7..3c8f706eb 100644 --- a/backend/internal/handler/user_account_handler.go +++ b/backend/internal/handler/user_account_handler.go @@ -124,6 +124,7 @@ type createUserAccountRequest struct { type importUserAccountCredentialsRequest struct { Contents []string `json:"contents" binding:"required"` Platform string `json:"platform" binding:"required,oneof=anthropic openai gemini antigravity grok"` + OpenAIAuthMode string `json:"openai_auth_mode" binding:"omitempty,oneof=oauth agent_identity"` AccountLevel string `json:"account_level"` ProxyID *int64 `json:"proxy_id"` ShareMode string `json:"share_mode" binding:"omitempty,oneof=private public"` @@ -213,6 +214,11 @@ const userOwnedDefaultConcurrency = 3 const userOwnedDefaultPriority = 1 const userAccountLevelVerifyLimitPerMinute = 5 +const ( + userOpenAIAuthModeOAuth = "oauth" + userOpenAIAuthModeAgentIdentity = "agent_identity" +) + type userOAuthProxyRequest struct { ProxyID *int64 `json:"proxy_id"` } @@ -462,6 +468,39 @@ func validateOpenAIImportTargetLevel(defaults importUserAccountCredentialsReques return targetLevel, nil } +func resolveUserOpenAICredentialImportMode( + req importUserAccountCredentialsRequest, + sources []service.AccountCredentialImportSource, +) (bool, error) { + declaredMode := strings.ToLower(strings.TrimSpace(req.OpenAIAuthMode)) + agentIdentityCount := 0 + for _, source := range sources { + if source.Kind == service.AccountCredentialImportKindOpenAIAgentIdentity { + agentIdentityCount++ + } + } + + if declaredMode == userOpenAIAuthModeAgentIdentity { + if req.Platform != service.PlatformOpenAI { + return false, infraerrors.BadRequest("OWNED_AGENT_IDENTITY_PLATFORM_INVALID", "Codex Agent Identity 仅支持 OpenAI 平台") + } + if agentIdentityCount != len(sources) { + return false, infraerrors.BadRequest("OWNED_AGENT_IDENTITY_CONTENT_INVALID", "Agent Identity 模式只接受 Agent Identity JSON 凭证") + } + return true, nil + } + if declaredMode == userOpenAIAuthModeOAuth && agentIdentityCount > 0 { + return false, infraerrors.BadRequest("OWNED_ACCOUNT_IMPORT_AUTH_MODE_MISMATCH", "导入凭证与所选 OpenAI 认证模式不一致") + } + if agentIdentityCount == 0 { + return false, nil + } + if req.Platform != service.PlatformOpenAI || agentIdentityCount != len(sources) { + return false, infraerrors.BadRequest("OWNED_ACCOUNT_IMPORT_AUTH_MODE_MIXED", "Agent Identity 凭证不能与其他认证凭证混合导入") + } + return true, nil +} + func userUnixSecondsToTime(value *int64) *time.Time { if value == nil || *value <= 0 { return nil @@ -1323,13 +1362,25 @@ func (h *UserAccountHandler) ImportCredentials(c *gin.Context) { response.BadRequest(c, "No importable account credentials found") return } + isAgentIdentityImport, err := resolveUserOpenAICredentialImportMode(req, sources) + if err != nil { + response.ErrorFrom(c, err) + return + } levelConfigs, err := h.openAIAccountLevelConfigs(c.Request.Context()) if err != nil { response.ErrorFrom(c, err) return } - normalizeUserCredentialImportTargetLevel(&req, levelConfigs) - if service.RequiresUserAccountOAuthProxyWithConfigs(req.Platform, service.AccountLevelUnknown, levelConfigs) { + if isAgentIdentityImport { + req.AccountLevel = service.AccountLevelUnknown + req.ProxyID = nil + req.ShareMode = service.AccountShareModePrivate + req.ExpiresAt = nil + } else { + normalizeUserCredentialImportTargetLevel(&req, levelConfigs) + } + if !isAgentIdentityImport && service.RequiresUserAccountOAuthProxyWithConfigs(req.Platform, service.AccountLevelUnknown, levelConfigs) { if !h.requireUserOAuthProxy(c, subject.UserID, req.ProxyID) { return } @@ -1351,7 +1402,7 @@ func (h *UserAccountHandler) ImportCredentials(c *gin.Context) { result.Errors = append(result.Errors, parseErrors...) for idx, source := range sources { - account, err := h.createOwnedAccountFromCredentialImportSource(c.Request.Context(), subject.UserID, source, req, idx+1) + outcome, err := h.createOwnedAccountFromCredentialImportSource(c.Request.Context(), subject.UserID, source, req, idx+1) if err != nil { result.Failed++ result.Errors = append(result.Errors, service.AccountCredentialImportError{ @@ -1362,8 +1413,12 @@ func (h *UserAccountHandler) ImportCredentials(c *gin.Context) { }) continue } - if account != nil { - result.Created++ + if outcome != nil && outcome.Account != nil { + if outcome.Updated { + result.Updated++ + } else { + result.Created++ + } } } result.Failed += len(parseErrors) @@ -1383,13 +1438,14 @@ func (h *UserAccountHandler) createOwnedAccountFromCredentialImportSource( source service.AccountCredentialImportSource, defaults importUserAccountCredentialsRequest, sequence int, -) (*service.Account, error) { +) (*service.OwnedAccountImportResult, error) { if err := validateCredentialImportTargetPlatform(defaults, source); err != nil { return nil, err } openAIAccountLevel := service.AccountLevelUnknown - if credentialImportSourceIsOpenAI(source) { + isAgentIdentity := source.Kind == service.AccountCredentialImportKindOpenAIAgentIdentity + if credentialImportSourceIsOpenAI(source) && !isAgentIdentity { levelConfigs, err := h.openAIAccountLevelConfigs(ctx) if err != nil { return nil, err @@ -1444,6 +1500,15 @@ func (h *UserAccountHandler) createOwnedAccountFromCredentialImportSource( if req.Name == "" { req.Name = fmt.Sprintf("OpenAI OAuth Account #%d", sequence) } + case service.AccountCredentialImportKindOpenAIAgentIdentity: + req.Platform = service.PlatformOpenAI + req.AccountLevel = service.AccountLevelUnknown + req.ShareMode = service.AccountShareModePrivate + req.ProxyID = nil + req.ExpiresAt = nil + if req.Name == "" { + req.Name = service.DeriveAccountCredentialImportName(req.Platform, req.Credentials, req.Extra, sequence) + } case service.AccountCredentialImportKindClaudeSessionKey: tokenInfo, err := h.oauthService.CookieAuth(ctx, &service.CookieAuthInput{ SessionKey: source.Token, @@ -1475,11 +1540,16 @@ func (h *UserAccountHandler) createOwnedAccountFromCredentialImportSource( if strings.TrimSpace(req.Name) == "" { return nil, fmt.Errorf("account name is required") } - account, err := h.accountService.ImportOwned(ctx, ownerUserID, req) + outcome, err := h.accountService.ImportOwnedWithResult(ctx, ownerUserID, req) + if err != nil { + return nil, err + } + account, err := h.activateOwnedPublicShareIfRequested(ctx, ownerUserID, outcome.Account) if err != nil { return nil, err } - return h.activateOwnedPublicShareIfRequested(ctx, ownerUserID, account) + outcome.Account = account + return outcome, nil } func (h *UserAccountHandler) Update(c *gin.Context) { @@ -1521,7 +1591,10 @@ func (h *UserAccountHandler) Update(c *gin.Context) { response.ErrorFrom(c, err) return } - if req.ShareMode != nil && service.NormalizeAccountShareMode(*req.ShareMode) == service.AccountShareModePublic { + requestedPublicShare := req.ShareMode != nil && service.NormalizeAccountShareMode(*req.ShareMode) == service.AccountShareModePublic + changedPublicAgentIdentityCredentials := req.Credentials != nil && account.IsOpenAIAgentIdentity() && + service.NormalizeAccountShareMode(account.ShareMode) == service.AccountShareModePublic + if requestedPublicShare || changedPublicAgentIdentityCredentials { account, err = h.activateOwnedPublicShareIfRequested(c.Request.Context(), subject.UserID, account) if err != nil { response.ErrorFrom(c, err) @@ -1842,14 +1915,18 @@ func (h *UserAccountHandler) BulkUpdate(c *gin.Context) { response.ErrorFrom(c, err) return } - if req.ShareMode != nil && service.NormalizeAccountShareMode(*req.ShareMode) == service.AccountShareModePublic { + requestedPublicShare := req.ShareMode != nil && service.NormalizeAccountShareMode(*req.ShareMode) == service.AccountShareModePublic + revalidatePublicAgentIdentity := len(req.Credentials) > 0 + if requestedPublicShare || revalidatePublicAgentIdentity { for i := range result.Results { entry := &result.Results[i] if !entry.Success { continue } account, err := h.accountService.GetOwnedByID(c.Request.Context(), subject.UserID, entry.AccountID) - if err == nil { + shouldActivate := account != nil && service.NormalizeAccountShareMode(account.ShareMode) == service.AccountShareModePublic && + (requestedPublicShare || account.IsOpenAIAgentIdentity()) + if err == nil && shouldActivate { _, err = h.activateOwnedPublicShareIfRequested(c.Request.Context(), subject.UserID, account) } if err != nil { diff --git a/backend/internal/handler/user_account_public_share_test.go b/backend/internal/handler/user_account_public_share_test.go index 5ce925d2b..624e19b90 100644 --- a/backend/internal/handler/user_account_public_share_test.go +++ b/backend/internal/handler/user_account_public_share_test.go @@ -1,11 +1,295 @@ package handler import ( + "bytes" + "context" + "crypto/ed25519" + "crypto/rand" + "crypto/x509" + "encoding/base64" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" "testing" + "github.com/Wei-Shaw/sub2api/internal/pkg/pagination" + "github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint" + 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" ) +const ( + userAgentIdentityPrivateGroupID int64 = 8201 + userAgentIdentityPublicGroupID int64 = 8202 +) + +type userAgentIdentityShareRepo struct { + service.AccountRepository + accounts map[int64]*service.Account +} + +func cloneUserAgentIdentityShareAccount(account *service.Account) *service.Account { + if account == nil { + return nil + } + clone := *account + clone.Credentials = make(map[string]any, len(account.Credentials)) + for key, value := range account.Credentials { + clone.Credentials[key] = value + } + clone.Extra = make(map[string]any, len(account.Extra)) + for key, value := range account.Extra { + clone.Extra[key] = value + } + clone.GroupIDs = append([]int64(nil), account.GroupIDs...) + if account.OwnerUserID != nil { + ownerUserID := *account.OwnerUserID + clone.OwnerUserID = &ownerUserID + } + return &clone +} + +func (r *userAgentIdentityShareRepo) GetByID(_ context.Context, accountID int64) (*service.Account, error) { + account := r.accounts[accountID] + if account == nil { + return nil, service.ErrAccountNotFound + } + return cloneUserAgentIdentityShareAccount(account), nil +} + +func (r *userAgentIdentityShareRepo) GetByIDs(_ context.Context, accountIDs []int64) ([]*service.Account, error) { + accounts := make([]*service.Account, 0, len(accountIDs)) + for _, accountID := range accountIDs { + if account := r.accounts[accountID]; account != nil { + accounts = append(accounts, cloneUserAgentIdentityShareAccount(account)) + } + } + return accounts, nil +} + +func (r *userAgentIdentityShareRepo) Update(_ context.Context, account *service.Account) error { + r.accounts[account.ID] = cloneUserAgentIdentityShareAccount(account) + return nil +} + +func (r *userAgentIdentityShareRepo) BindGroups(_ context.Context, accountID int64, groupIDs []int64) error { + account := r.accounts[accountID] + if account == nil { + return service.ErrAccountNotFound + } + account.GroupIDs = append([]int64(nil), groupIDs...) + return nil +} + +func (r *userAgentIdentityShareRepo) IsAccountShareModeListingAccount(context.Context, int64) (bool, error) { + return false, nil +} + +func (r *userAgentIdentityShareRepo) ListOwnedWithFilters( + _ context.Context, + ownerUserID int64, + params pagination.PaginationParams, + platform, accountType, _ string, + _ string, + _, _ int64, + _ string, +) ([]service.Account, *pagination.PaginationResult, error) { + accounts := make([]service.Account, 0, len(r.accounts)) + for _, account := range r.accounts { + if account.OwnerUserID == nil || *account.OwnerUserID != ownerUserID { + continue + } + if platform != "" && account.Platform != platform { + continue + } + if accountType != "" && account.Type != accountType { + continue + } + accounts = append(accounts, *cloneUserAgentIdentityShareAccount(account)) + } + total := int64(len(accounts)) + start := params.Offset() + if start >= len(accounts) { + return []service.Account{}, &pagination.PaginationResult{Total: total}, nil + } + end := start + params.Limit() + if end > len(accounts) { + end = len(accounts) + } + return accounts[start:end], &pagination.PaginationResult{Total: total}, nil +} + +type userAgentIdentityPrivateGroupProvisioner struct{} + +func (userAgentIdentityPrivateGroupProvisioner) ProvisionUserPrivateGroups(context.Context, int64) error { + return nil +} + +func (userAgentIdentityPrivateGroupProvisioner) GetActiveUserPrivateGroup(context.Context, int64, string) (*service.Group, error) { + return &service.Group{ + ID: userAgentIdentityPrivateGroupID, + Name: "OpenAI private", + Platform: service.PlatformOpenAI, + Status: service.StatusActive, + }, nil +} + +type userAgentIdentityPublicGroupRepo struct { + service.GroupRepository +} + +func (userAgentIdentityPublicGroupRepo) ListActiveByPlatform(_ context.Context, platform string) ([]service.Group, error) { + if platform != service.PlatformOpenAI { + return nil, nil + } + return []service.Group{{ + ID: userAgentIdentityPublicGroupID, + Name: "TEAM shared pool", + Platform: service.PlatformOpenAI, + Status: service.StatusActive, + Scope: service.GroupScopePublic, + RequiredAccountLevel: service.AccountLevelTeam, + }}, nil +} + +func (userAgentIdentityPublicGroupRepo) IsModeGroup(context.Context, int64) (bool, error) { + return false, nil +} + +type userAgentIdentitySharePolicyRepo struct { + service.AccountSharePolicyRepository +} + +func (userAgentIdentitySharePolicyRepo) ResolveEnabledAccountSharePolicy(context.Context, int64, *int64, string, *int64) (*service.AccountSharePolicy, error) { + return &service.AccountSharePolicy{ID: 1, Enabled: true, OwnerShareRatio: 0.7}, nil +} + +type userAgentIdentityValidationUpstream struct { + statusCode int + body string + calls int + lastAuthorization string +} + +func (u *userAgentIdentityValidationUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) { + return u.response(req), nil +} + +func (u *userAgentIdentityValidationUpstream) DoWithTLS(req *http.Request, _ string, _ int64, _ int, _ *tlsfingerprint.Profile) (*http.Response, error) { + return u.response(req), nil +} + +func (u *userAgentIdentityValidationUpstream) response(req *http.Request) *http.Response { + u.calls++ + u.lastAuthorization = req.Header.Get("Authorization") + statusCode := u.statusCode + if statusCode == 0 { + statusCode = http.StatusOK + } + body := u.body + if body == "" && statusCode == http.StatusOK { + body = "data: {\"type\":\"response.completed\"}\n\n" + } + return &http.Response{ + StatusCode: statusCode, + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader(body)), + } +} + +type userAgentIdentityWSInvalidationRecorder struct { + accountIDs []int64 +} + +func (r *userAgentIdentityWSInvalidationRecorder) InvalidateAgentIdentityWSConnections(accountID int64) { + r.accountIDs = append(r.accountIDs, accountID) +} + +func userAgentIdentityCredentials(t *testing.T, runtimeID, taskID string) map[string]any { + t.Helper() + _, privateKey, err := ed25519.GenerateKey(rand.Reader) + require.NoError(t, err) + der, err := x509.MarshalPKCS8PrivateKey(privateKey) + require.NoError(t, err) + return map[string]any{ + "auth_mode": service.OpenAIAuthModeAgentIdentity, + "agent_runtime_id": runtimeID, + "agent_private_key": base64.StdEncoding.EncodeToString(der), + "task_id": taskID, + "chatgpt_account_id": "team-handler-test", + "chatgpt_user_id": "member-handler-test", + "plan_type": service.AccountLevelTeam, + } +} + +func newUserAgentIdentityShareAccount(t *testing.T, ownerUserID int64, shareMode, shareStatus string) *service.Account { + t.Helper() + groupIDs := []int64{userAgentIdentityPrivateGroupID} + if shareMode == service.AccountShareModePublic && shareStatus == service.AccountShareStatusApproved { + groupIDs = append(groupIDs, userAgentIdentityPublicGroupID) + } + return &service.Account{ + ID: 1, + Name: "Agent Identity", + OwnerUserID: &ownerUserID, + Platform: service.PlatformOpenAI, + Type: service.AccountTypeOAuth, + AccountLevel: service.AccountLevelTeam, + Credentials: userAgentIdentityCredentials(t, "runtime-old", "task-old"), + Extra: map[string]any{}, + ShareMode: shareMode, + ShareStatus: shareStatus, + Concurrency: 3, + Priority: 5, + Status: service.StatusActive, + Schedulable: true, + GroupIDs: groupIDs, + } +} + +func newUserAgentIdentityShareHandler( + t *testing.T, + account *service.Account, + upstreamStatus int, + upstreamBody string, +) (*UserAccountHandler, *userAgentIdentityShareRepo, *userAgentIdentityValidationUpstream, *userAgentIdentityWSInvalidationRecorder) { + t.Helper() + repo := &userAgentIdentityShareRepo{ + accounts: map[int64]*service.Account{account.ID: cloneUserAgentIdentityShareAccount(account)}, + } + invalidator := &userAgentIdentityWSInvalidationRecorder{} + invalidatorProxy := service.NewAgentIdentityWSInvalidatorProxy() + invalidatorProxy.SetTarget(invalidator) + accountService := service.NewAccountService(repo, userAgentIdentityPublicGroupRepo{}, nil, nil, nil) + accountService.SetUserPrivateGroupProvisioner(userAgentIdentityPrivateGroupProvisioner{}) + accountService.SetAccountSharePolicyRepository(userAgentIdentitySharePolicyRepo{}) + accountService.SetAgentIdentityWSInvalidator(invalidatorProxy) + upstream := &userAgentIdentityValidationUpstream{statusCode: upstreamStatus, body: upstreamBody} + accountTestService := service.NewAccountTestService(repo, nil, nil, nil, upstream, nil, nil, nil, invalidatorProxy) + handler := NewUserAccountHandler(accountService, nil, accountTestService, nil, nil, nil, nil, nil, nil, nil) + return handler, repo, upstream, invalidator +} + +func runUserAgentIdentityUpdateRequest(t *testing.T, handler *UserAccountHandler, ownerUserID int64, body any) *httptest.ResponseRecorder { + t.Helper() + payload, err := json.Marshal(body) + require.NoError(t, err) + router := gin.New() + router.PUT("/accounts/:id", func(c *gin.Context) { + c.Set(string(middleware2.ContextKeyUser), middleware2.AuthSubject{UserID: ownerUserID}) + handler.Update(c) + }) + request := httptest.NewRequest(http.MethodPut, "/accounts/1", bytes.NewReader(payload)) + request.Header.Set("Content-Type", "application/json") + recorder := httptest.NewRecorder() + router.ServeHTTP(recorder, request) + return recorder +} + func TestIsOpenAIUsageLimitReachedValidationError(t *testing.T) { require.True(t, isOpenAIUsageLimitReachedValidationError(`API returned 429: {"error":{"type":"usage_limit_reached","message":"The usage limit has been reached"}}`)) require.True(t, isOpenAIUsageLimitReachedValidationError(`API returned 429: {"error": {"type": "usage_limit_reached"}}`)) @@ -13,3 +297,78 @@ func TestIsOpenAIUsageLimitReachedValidationError(t *testing.T) { require.False(t, isOpenAIUsageLimitReachedValidationError(`API returned 401: {"error":{"type":"usage_limit_reached"}}`)) require.False(t, isOpenAIUsageLimitReachedValidationError(`Request failed: dial tcp timeout`)) } + +func TestUserAccountHandlerUpdateAgentIdentityPrivateToPublicRevalidates(t *testing.T) { + gin.SetMode(gin.TestMode) + ownerUserID := int64(101) + + t.Run("approves after successful connection test", func(t *testing.T) { + account := newUserAgentIdentityShareAccount(t, ownerUserID, service.AccountShareModePrivate, service.AccountShareStatusApproved) + handler, repo, upstream, _ := newUserAgentIdentityShareHandler(t, account, http.StatusOK, "") + + recorder := runUserAgentIdentityUpdateRequest(t, handler, ownerUserID, map[string]any{"share_mode": service.AccountShareModePublic}) + + require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String()) + require.Equal(t, 1, upstream.calls) + require.NotEmpty(t, upstream.lastAuthorization) + stored := repo.accounts[account.ID] + require.Equal(t, service.AccountShareModePublic, stored.ShareMode) + require.Equal(t, service.AccountShareStatusApproved, stored.ShareStatus) + require.Empty(t, stored.ErrorMessage) + require.Equal(t, []int64{userAgentIdentityPrivateGroupID, userAgentIdentityPublicGroupID}, stored.GroupIDs) + }) + + t.Run("keeps account pending after failed connection test", func(t *testing.T) { + account := newUserAgentIdentityShareAccount(t, ownerUserID, service.AccountShareModePrivate, service.AccountShareStatusApproved) + handler, repo, upstream, _ := newUserAgentIdentityShareHandler(t, account, http.StatusServiceUnavailable, `{"error":"upstream unavailable"}`) + + recorder := runUserAgentIdentityUpdateRequest(t, handler, ownerUserID, map[string]any{"share_mode": service.AccountShareModePublic}) + + require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String()) + require.Equal(t, 1, upstream.calls) + stored := repo.accounts[account.ID] + require.Equal(t, service.AccountShareModePublic, stored.ShareMode) + require.Equal(t, service.AccountShareStatusPending, stored.ShareStatus) + require.Contains(t, stored.ErrorMessage, "API returned 503") + require.Equal(t, []int64{userAgentIdentityPrivateGroupID}, stored.GroupIDs) + }) +} + +func TestUserAccountHandlerUpdateApprovedPublicAgentIdentityCredentialsRevalidates(t *testing.T) { + gin.SetMode(gin.TestMode) + ownerUserID := int64(101) + account := newUserAgentIdentityShareAccount(t, ownerUserID, service.AccountShareModePublic, service.AccountShareStatusApproved) + handler, repo, upstream, invalidator := newUserAgentIdentityShareHandler(t, account, http.StatusOK, "") + newCredentials := userAgentIdentityCredentials(t, "runtime-new", "task-new") + + recorder := runUserAgentIdentityUpdateRequest(t, handler, ownerUserID, map[string]any{"credentials": newCredentials}) + + require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String()) + require.Equal(t, 1, upstream.calls, "changed public credentials must trigger a fresh connection test") + require.Equal(t, []int64{account.ID}, invalidator.accountIDs) + stored := repo.accounts[account.ID] + require.Equal(t, "runtime-new", stored.GetCredential("agent_runtime_id")) + require.Equal(t, service.AccountShareModePublic, stored.ShareMode) + require.Equal(t, service.AccountShareStatusApproved, stored.ShareStatus) + require.Equal(t, []int64{userAgentIdentityPrivateGroupID, userAgentIdentityPublicGroupID}, stored.GroupIDs) +} + +func TestUserAccountHandlerSetPublicShareExecutorValidatesAgentIdentityBeforeApproval(t *testing.T) { + ownerUserID := int64(101) + account := newUserAgentIdentityShareAccount(t, ownerUserID, service.AccountShareModePrivate, service.AccountShareStatusApproved) + handler, repo, upstream, _ := newUserAgentIdentityShareHandler(t, account, http.StatusOK, "") + task := &service.AccountBatchTask{ + Operation: service.AccountBatchTaskOperationUserSetPublicShare, + OwnerUserID: &ownerUserID, + } + + result, err := handler.executeUserSetPublicShareTaskItem(context.Background(), task, service.AccountBatchTaskItem{AccountID: account.ID}) + + require.NoError(t, err) + require.Equal(t, 1, upstream.calls, "the async executor must not bypass public-share validation") + require.Equal(t, service.AccountShareModePublic, result["share_mode"]) + require.Equal(t, service.AccountShareStatusApproved, result["share_status"]) + stored := repo.accounts[account.ID] + require.Equal(t, service.AccountShareStatusApproved, stored.ShareStatus) + require.Equal(t, []int64{userAgentIdentityPrivateGroupID, userAgentIdentityPublicGroupID}, stored.GroupIDs) +} diff --git a/backend/internal/handler/wire.go b/backend/internal/handler/wire.go index b57fcad88..0fc7c7177 100644 --- a/backend/internal/handler/wire.go +++ b/backend/internal/handler/wire.go @@ -1,6 +1,7 @@ package handler import ( + "github.com/Wei-Shaw/sub2api/internal/config" "github.com/Wei-Shaw/sub2api/internal/handler/admin" "github.com/Wei-Shaw/sub2api/internal/service" @@ -102,6 +103,33 @@ func ProvideSettingHandler(settingService *service.SettingService, buildInfo Bui return NewSettingHandler(settingService, buildInfo.Version) } +func ProvideOpenAIGatewayHandler( + gatewayService *service.OpenAIGatewayService, + concurrencyService *service.ConcurrencyService, + billingCacheService *service.BillingCacheService, + apiKeyService *service.APIKeyService, + usageRecordWorkerPool *service.UsageRecordWorkerPool, + errorPassthroughService *service.ErrorPassthroughService, + contentModerationService *service.ContentModerationService, + userModerationService *service.UserContentModerationService, + grokQuotaService *service.GrokQuotaService, + cfg *config.Config, +) *OpenAIGatewayHandler { + h := NewOpenAIGatewayHandler( + gatewayService, + concurrencyService, + billingCacheService, + apiKeyService, + usageRecordWorkerPool, + errorPassthroughService, + contentModerationService, + userModerationService, + cfg, + ) + h.grokMediaEligibilityProber = grokQuotaService + return h +} + func ProvideAdminAccountHandler( adminService service.AdminService, accountService *service.AccountService, @@ -255,7 +283,7 @@ var ProviderSet = wire.NewSet( NewConversationHandler, NewChannelMonitorUserHandler, NewGatewayHandler, - NewOpenAIGatewayHandler, + ProvideOpenAIGatewayHandler, NewTotpHandler, ProvideSettingHandler, NewPaymentHandler, diff --git a/backend/internal/pkg/apicompat/anthropic_responses_test.go b/backend/internal/pkg/apicompat/anthropic_responses_test.go index 393d885b1..4e7e3e2a1 100644 --- a/backend/internal/pkg/apicompat/anthropic_responses_test.go +++ b/backend/internal/pkg/apicompat/anthropic_responses_test.go @@ -181,6 +181,20 @@ func TestResponsesToAnthropic_TextOnly(t *testing.T) { assert.Equal(t, 5, anth.Usage.OutputTokens) } +func TestAnthropicResponseMarshalJSONUsesNullForEmptyStopReason(t *testing.T) { + payload, err := json.Marshal(AnthropicResponse{ + ID: "msg_start", + Type: "message", + Role: "assistant", + Content: []AnthropicContentBlock{}, + Model: "claude-test", + }) + + require.NoError(t, err) + require.Contains(t, string(payload), `"stop_reason":null`) + require.NotContains(t, string(payload), `"stop_reason":""`) +} + func TestResponsesToAnthropic_CachedTokensUseAnthropicInputSemantics(t *testing.T) { resp := &ResponsesResponse{ ID: "resp_cached", diff --git a/backend/internal/pkg/apicompat/anthropic_to_responses_response.go b/backend/internal/pkg/apicompat/anthropic_to_responses_response.go index 892663f19..f3fc0e833 100644 --- a/backend/internal/pkg/apicompat/anthropic_to_responses_response.go +++ b/backend/internal/pkg/apicompat/anthropic_to_responses_response.go @@ -149,11 +149,24 @@ type AnthropicEventToResponsesState struct { // For message output: accumulate text parts ContentIndex int + // TextAccum accumulates the current text part so output_text.done and + // content_part.done can carry the complete text instead of only deltas. + TextAccum string // For function_call: track per-output info CurrentCallID string CurrentName string + // Accumulated payload for the currently open output item. + CurrentContent []ResponsesContentPart + CurrentArgs string + CurrentSummary string + + // Outputs contains every closed item and is emitted on the terminal event. + // OpenAI SDK accumulators read response.completed directly when producing + // get_final_response(), so the terminal response cannot use an empty list. + Outputs []ResponsesOutput + // Usage from Anthropic stream events. InputTokens uses Anthropic semantics and excludes cache tokens. InputTokens int OutputTokens int @@ -288,6 +301,17 @@ func anthToResHandleContentBlockStart(evt *AnthropicStreamEvent, state *Anthropi })) } + // The message item starts with an empty content array. OpenAI SDK stream + // accumulation requires content_part.added before the first text delta so + // it can create output.content[content_index]. + events = append(events, makeResponsesEvent(state, "response.content_part.added", &ResponsesStreamEvent{ + OutputIndex: state.OutputIndex, + ContentIndex: state.ContentIndex, + ItemID: state.CurrentItemID, + Part: &ResponsesContentPart{Type: "output_text", Text: ""}, + })) + state.TextAccum = "" + case "tool_use": // Close previous item if any events = append(events, closeCurrentResponsesItem(state)...) @@ -322,6 +346,7 @@ func anthToResHandleContentBlockDelta(evt *AnthropicStreamEvent, state *Anthropi if evt.Delta.Text == "" { return nil } + state.TextAccum += evt.Delta.Text return []ResponsesStreamEvent{makeResponsesEvent(state, "response.output_text.delta", &ResponsesStreamEvent{ OutputIndex: state.OutputIndex, ContentIndex: state.ContentIndex, @@ -333,6 +358,7 @@ func anthToResHandleContentBlockDelta(evt *AnthropicStreamEvent, state *Anthropi if evt.Delta.Thinking == "" { return nil } + state.CurrentSummary += evt.Delta.Thinking return []ResponsesStreamEvent{makeResponsesEvent(state, "response.reasoning_summary_text.delta", &ResponsesStreamEvent{ OutputIndex: state.OutputIndex, SummaryIndex: 0, @@ -344,6 +370,7 @@ func anthToResHandleContentBlockDelta(evt *AnthropicStreamEvent, state *Anthropi if evt.Delta.PartialJSON == "" { return nil } + state.CurrentArgs += evt.Delta.PartialJSON return []ResponsesStreamEvent{makeResponsesEvent(state, "response.function_call_arguments.delta", &ResponsesStreamEvent{ OutputIndex: state.OutputIndex, Delta: evt.Delta.PartialJSON, @@ -388,14 +415,29 @@ func anthToResHandleContentBlockStop(evt *AnthropicStreamEvent, state *Anthropic return events case "message": - // Emit output_text.done (text block is done, but message item stays open for potential more blocks) - return []ResponsesStreamEvent{ + // Complete this content part while keeping the message item open for a + // possible following text block. Done events carry the full text. + text := state.TextAccum + state.TextAccum = "" + state.CurrentContent = append(state.CurrentContent, ResponsesContentPart{Type: "output_text", Text: text}) + events := []ResponsesStreamEvent{ makeResponsesEvent(state, "response.output_text.done", &ResponsesStreamEvent{ OutputIndex: state.OutputIndex, ContentIndex: state.ContentIndex, ItemID: state.CurrentItemID, + Text: text, + }), + makeResponsesEvent(state, "response.content_part.done", &ResponsesStreamEvent{ + OutputIndex: state.OutputIndex, + ContentIndex: state.ContentIndex, + ItemID: state.CurrentItemID, + Part: &ResponsesContentPart{Type: "output_text", Text: text}, }), } + // Anthropic may emit multiple text blocks in one assistant message. Each + // Responses content part needs a distinct index for SDK accumulation. + state.ContentIndex++ + return events } return nil @@ -452,24 +494,44 @@ func closeCurrentResponsesItem(state *AnthropicEventToResponsesState) []Response return nil } - itemType := state.CurrentItemType - itemID := state.CurrentItemID + item := ResponsesOutput{ + Type: state.CurrentItemType, + ID: state.CurrentItemID, + Status: "completed", + } + switch state.CurrentItemType { + case "message": + item.Role = "assistant" + item.Content = state.CurrentContent + case "function_call": + item.CallID = state.CurrentCallID + item.Name = state.CurrentName + item.Arguments = state.CurrentArgs + if item.Arguments == "" { + item.Arguments = "{}" + } + case "reasoning": + if state.CurrentSummary != "" { + item.Summary = []ResponsesSummary{{Type: "summary_text", Text: state.CurrentSummary}} + } + } + state.Outputs = append(state.Outputs, item) // Reset state.CurrentItemType = "" state.CurrentItemID = "" state.CurrentCallID = "" state.CurrentName = "" + state.CurrentContent = nil + state.CurrentArgs = "" + state.CurrentSummary = "" + state.TextAccum = "" state.OutputIndex++ state.ContentIndex = 0 return []ResponsesStreamEvent{makeResponsesEvent(state, "response.output_item.done", &ResponsesStreamEvent{ OutputIndex: state.OutputIndex - 1, // Use the index before increment - Item: &ResponsesOutput{ - Type: itemType, - ID: itemID, - Status: "completed", - }, + Item: &item, })} } @@ -514,6 +576,10 @@ func makeResponsesCompletedEvent( if status == "incomplete" { eventType = "response.incomplete" } + outputs := state.Outputs + if outputs == nil { + outputs = []ResponsesOutput{} + } return ResponsesStreamEvent{ Type: eventType, @@ -523,7 +589,7 @@ func makeResponsesCompletedEvent( Object: "response", Model: state.Model, Status: status, - Output: []ResponsesOutput{}, // Simplified; full output tracking would add complexity + Output: outputs, Usage: usage, IncompleteDetails: incompleteDetails, }, diff --git a/backend/internal/pkg/apicompat/anthropic_to_responses_stream_test.go b/backend/internal/pkg/apicompat/anthropic_to_responses_stream_test.go new file mode 100644 index 000000000..307426d5a --- /dev/null +++ b/backend/internal/pkg/apicompat/anthropic_to_responses_stream_test.go @@ -0,0 +1,184 @@ +package apicompat + +import "testing" + +func TestAnthropicEventToResponsesTextEmitsOrderedContentPartEvents(t *testing.T) { + state := NewAnthropicEventToResponsesState() + state.Model = "claude-sonnet-4-5" + + var events []ResponsesStreamEvent + feed := func(event *AnthropicStreamEvent) { + events = append(events, AnthropicEventToResponsesEvents(event, state)...) + } + + index := 0 + feed(&AnthropicStreamEvent{Type: "message_start", Message: &AnthropicResponse{ID: "msg_1", Model: state.Model}}) + feed(&AnthropicStreamEvent{Type: "content_block_start", Index: &index, ContentBlock: &AnthropicContentBlock{Type: "text"}}) + feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &index, Delta: &AnthropicDelta{Type: "text_delta", Text: "Hel"}}) + feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &index, Delta: &AnthropicDelta{Type: "text_delta", Text: "lo"}}) + feed(&AnthropicStreamEvent{Type: "content_block_stop", Index: &index}) + feed(&AnthropicStreamEvent{Type: "message_stop"}) + + position := func(eventType string) int { + for i := range events { + if events[i].Type == eventType { + return i + } + } + return -1 + } + partAdded := position("response.content_part.added") + firstDelta := position("response.output_text.delta") + textDone := position("response.output_text.done") + partDone := position("response.content_part.done") + if partAdded < 0 || firstDelta < 0 || textDone < 0 || partDone < 0 { + t.Fatalf("missing required content events: %+v", eventTypes(events)) + } + if !(partAdded < firstDelta && firstDelta < textDone && textDone < partDone) { + t.Fatalf("invalid content event order: %+v", eventTypes(events)) + } + if events[textDone].Text != "Hello" { + t.Fatalf("output_text.done text = %q, want Hello", events[textDone].Text) + } + if events[partDone].Part == nil || events[partDone].Part.Text != "Hello" { + t.Fatalf("content_part.done part = %+v, want full text", events[partDone].Part) + } +} + +func TestAnthropicEventToResponsesMultipleTextPartsUseDistinctIndexes(t *testing.T) { + state := NewAnthropicEventToResponsesState() + var events []ResponsesStreamEvent + feed := func(event *AnthropicStreamEvent) { + events = append(events, AnthropicEventToResponsesEvents(event, state)...) + } + + feed(&AnthropicStreamEvent{Type: "message_start", Message: &AnthropicResponse{ID: "msg_multi"}}) + for index, text := range []string{"first", "second"} { + blockIndex := index + feed(&AnthropicStreamEvent{Type: "content_block_start", Index: &blockIndex, ContentBlock: &AnthropicContentBlock{Type: "text"}}) + feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &blockIndex, Delta: &AnthropicDelta{Type: "text_delta", Text: text}}) + feed(&AnthropicStreamEvent{Type: "content_block_stop", Index: &blockIndex}) + } + feed(&AnthropicStreamEvent{Type: "message_stop"}) + + var addedIndexes []int + var completed *ResponsesStreamEvent + for i := range events { + switch events[i].Type { + case "response.content_part.added": + addedIndexes = append(addedIndexes, events[i].ContentIndex) + case "response.completed": + completed = &events[i] + } + } + if len(addedIndexes) != 2 || addedIndexes[0] != 0 || addedIndexes[1] != 1 { + t.Fatalf("content_part.added indexes = %v, want [0 1]", addedIndexes) + } + if completed == nil || completed.Response == nil || len(completed.Response.Output) != 1 { + t.Fatalf("terminal output missing: %+v", completed) + } + content := completed.Response.Output[0].Content + if len(content) != 2 || content[0].Text != "first" || content[1].Text != "second" { + t.Fatalf("terminal message content = %+v", content) + } +} + +func TestAnthropicEventToResponsesCompletedCarriesFullTextOutput(t *testing.T) { + state := NewAnthropicEventToResponsesState() + state.Model = "claude-sonnet-4-5" + var events []ResponsesStreamEvent + feed := func(event *AnthropicStreamEvent) { + events = append(events, AnthropicEventToResponsesEvents(event, state)...) + } + + index := 0 + feed(&AnthropicStreamEvent{Type: "message_start", Message: &AnthropicResponse{ID: "msg_1"}}) + feed(&AnthropicStreamEvent{Type: "content_block_start", Index: &index, ContentBlock: &AnthropicContentBlock{Type: "text"}}) + feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &index, Delta: &AnthropicDelta{Type: "text_delta", Text: "4826"}}) + feed(&AnthropicStreamEvent{Type: "content_block_stop", Index: &index}) + feed(&AnthropicStreamEvent{Type: "message_stop"}) + + completed := terminalResponseEvent(events) + if completed == nil || completed.Response == nil || len(completed.Response.Output) != 1 { + t.Fatalf("response.completed carries no output") + } + message := completed.Response.Output[0] + if message.Type != "message" || message.Role != "assistant" || len(message.Content) != 1 { + t.Fatalf("terminal output = %+v, want completed assistant message", message) + } + if message.Content[0].Text != "4826" { + t.Fatalf("terminal output text = %q, want 4826", message.Content[0].Text) + } +} + +func TestAnthropicEventToResponsesToolCallCarriesArgumentsInDoneAndCompleted(t *testing.T) { + state := NewAnthropicEventToResponsesState() + var events []ResponsesStreamEvent + feed := func(event *AnthropicStreamEvent) { + events = append(events, AnthropicEventToResponsesEvents(event, state)...) + } + + index := 0 + feed(&AnthropicStreamEvent{Type: "message_start", Message: &AnthropicResponse{ID: "msg_tool"}}) + feed(&AnthropicStreamEvent{Type: "content_block_start", Index: &index, ContentBlock: &AnthropicContentBlock{ + Type: "tool_use", ID: "toolu_1", Name: "get_weather", + }}) + feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &index, Delta: &AnthropicDelta{Type: "input_json_delta", PartialJSON: `{"city":`}}) + feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &index, Delta: &AnthropicDelta{Type: "input_json_delta", PartialJSON: `"SH"}`}}) + feed(&AnthropicStreamEvent{Type: "content_block_stop", Index: &index}) + feed(&AnthropicStreamEvent{Type: "message_stop"}) + + var doneItem *ResponsesOutput + for i := range events { + if events[i].Type == "response.output_item.done" { + doneItem = events[i].Item + } + } + completed := terminalResponseEvent(events) + if doneItem == nil || completed == nil || completed.Response == nil || len(completed.Response.Output) != 1 { + t.Fatalf("tool call output missing from done or terminal event") + } + for label, call := range map[string]ResponsesOutput{ + "output_item.done": *doneItem, + "response.completed": completed.Response.Output[0], + } { + if call.Type != "function_call" || call.Name != "get_weather" || call.Arguments != `{"city":"SH"}` { + t.Fatalf("%s tool call = %+v", label, call) + } + } +} + +func TestFinalizeAnthropicResponsesStreamCarriesAccumulatedOutput(t *testing.T) { + state := NewAnthropicEventToResponsesState() + index := 0 + AnthropicEventToResponsesEvents(&AnthropicStreamEvent{Type: "message_start", Message: &AnthropicResponse{ID: "msg_truncated"}}, state) + AnthropicEventToResponsesEvents(&AnthropicStreamEvent{Type: "content_block_start", Index: &index, ContentBlock: &AnthropicContentBlock{Type: "text"}}, state) + AnthropicEventToResponsesEvents(&AnthropicStreamEvent{Type: "content_block_delta", Index: &index, Delta: &AnthropicDelta{Type: "text_delta", Text: "partial"}}, state) + AnthropicEventToResponsesEvents(&AnthropicStreamEvent{Type: "content_block_stop", Index: &index}, state) + + events := FinalizeAnthropicResponsesStream(state) + completed := terminalResponseEvent(events) + if completed == nil || completed.Response == nil || len(completed.Response.Output) != 1 { + t.Fatalf("synthetic terminal event carries no output: %+v", events) + } + if got := completed.Response.Output[0].Content[0].Text; got != "partial" { + t.Fatalf("synthetic terminal output text = %q, want partial", got) + } +} + +func eventTypes(events []ResponsesStreamEvent) []string { + types := make([]string, 0, len(events)) + for i := range events { + types = append(types, events[i].Type) + } + return types +} + +func terminalResponseEvent(events []ResponsesStreamEvent) *ResponsesStreamEvent { + for i := range events { + if events[i].Type == "response.completed" || events[i].Type == "response.incomplete" { + return &events[i] + } + } + return nil +} diff --git a/backend/internal/pkg/apicompat/types.go b/backend/internal/pkg/apicompat/types.go index fcaabdc82..f608b1efb 100644 --- a/backend/internal/pkg/apicompat/types.go +++ b/backend/internal/pkg/apicompat/types.go @@ -108,6 +108,24 @@ type AnthropicResponse struct { Usage AnthropicUsage `json:"usage"` } +// MarshalJSON keeps the internal string representation while matching the +// Anthropic streaming contract: message_start must carry stop_reason:null, +// not an empty string. Final responses continue to emit their concrete reason. +func (r AnthropicResponse) MarshalJSON() ([]byte, error) { + type responseAlias AnthropicResponse + var stopReason *string + if r.StopReason != "" { + stopReason = &r.StopReason + } + return json.Marshal(struct { + responseAlias + StopReason *string `json:"stop_reason"` + }{ + responseAlias: responseAlias(r), + StopReason: stopReason, + }) +} + // AnthropicUsage holds token counts in Anthropic format. type AnthropicUsage struct { InputTokens int `json:"input_tokens"` diff --git a/backend/internal/pkg/claude/constants.go b/backend/internal/pkg/claude/constants.go index 011444bb4..aec4ad023 100644 --- a/backend/internal/pkg/claude/constants.go +++ b/backend/internal/pkg/claude/constants.go @@ -48,7 +48,8 @@ const MessageBetaHeaderWithTools = BetaClaudeCode + "," + BetaOAuth + "," + Beta // CountTokensBetaHeader count_tokens 请求使用的 anthropic-beta header const CountTokensBetaHeader = BetaClaudeCode + "," + BetaOAuth + "," + BetaInterleavedThinking + "," + BetaTokenCounting -// HaikuBetaHeader Haiku 模型使用的 anthropic-beta header（不需要 claude-code beta） +// HaikuBetaHeader Haiku 模型在 OAuth 真实客户端透传路径上的默认 anthropic-beta header。 +// OAuth mimic 路径统一使用 FullClaudeCodeMimicryBetas。 const HaikuBetaHeader = BetaOAuth + "," + BetaInterleavedThinking // APIKeyBetaHeader API-key 账号建议使用的 anthropic-beta header（不包含 oauth） @@ -72,8 +73,8 @@ const CLICurrentVersion = "2.1.92" // 顺序与真实 CLI 抓包一致。 // // 使用建议： -// - OAuth 账号 + 非 haiku：追加这整份列表，再按需保留 client 带来的 beta。 -// - OAuth 账号 + haiku：Anthropic 对 haiku 不做 third-party 判定，使用 HaikuBetaHeader 即可。 +// - OAuth mimic：所有模型（包括 Haiku）都使用这整份列表。 +// - OAuth 真实客户端透传：保留客户端 beta；未提供时使用模型对应默认值。 // - API-key 账号：不要使用本函数，参见 APIKeyBetaHeader。 func FullClaudeCodeMimicryBetas() []string { return []string{ diff --git a/backend/internal/pkg/ip/ip.go b/backend/internal/pkg/ip/ip.go index f6f77c86e..ed1950eaf 100644 --- a/backend/internal/pkg/ip/ip.go +++ b/backend/internal/pkg/ip/ip.go @@ -8,46 +8,25 @@ import ( "github.com/gin-gonic/gin" ) -// GetClientIP 从 Gin Context 中提取客户端真实 IP 地址。 -// 按以下优先级检查 Header： -// 1. CF-Connecting-IP (Cloudflare) -// 2. X-Real-IP (Nginx) -// 3. X-Forwarded-For (取第一个非私有 IP) -// 4. c.ClientIP() (Gin 内置方法) +// GetClientIP 从 Gin 的可信代理链中提取客户端 IP。 +// +// Deprecated: 安全敏感调用应显式使用 GetSecurityClientIP。保留该入口是为了让 +// 网关日志和会话种子等既有调用同样遵循可信代理规则，而不是直接信任原始请求头。 func GetClientIP(c *gin.Context) string { - // 1. Cloudflare - if ip := c.GetHeader("CF-Connecting-IP"); ip != "" { - return normalizeIP(ip) - } - - // 2. Nginx X-Real-IP - if ip := c.GetHeader("X-Real-IP"); ip != "" { - return normalizeIP(ip) - } - - // 3. X-Forwarded-For (多个 IP 时取第一个公网 IP) - if xff := c.GetHeader("X-Forwarded-For"); xff != "" { - ips := strings.Split(xff, ",") - for _, ip := range ips { - ip = strings.TrimSpace(ip) - if ip != "" && !isPrivateIP(ip) { - return normalizeIP(ip) - } - } - // 如果都是私有 IP，返回第一个 - if len(ips) > 0 { - return normalizeIP(strings.TrimSpace(ips[0])) - } - } - - // 4. Gin 内置方法 - return normalizeIP(c.ClientIP()) + return GetSecurityClientIP(c) } // GetTrustedClientIP 从 Gin 的可信代理解析链提取客户端 IP。 // 该方法依赖 gin.Engine.SetTrustedProxies 配置，不会优先直接信任原始转发头值。 // 适用于 ACL / 风控等安全敏感场景。 func GetTrustedClientIP(c *gin.Context) string { + return GetSecurityClientIP(c) +} + +// GetSecurityClientIP 是登录、注册、风控、ACL 与审计等安全敏感路径的统一入口。 +// Gin 仅在请求直连来源命中 server.trusted_proxies 时解析配置的转发头；否则返回 +// Request.RemoteAddr 中的直连地址。 +func GetSecurityClientIP(c *gin.Context) string { if c == nil { return "" } @@ -64,9 +43,6 @@ func normalizeIP(ip string) string { return ip } -// privateNets 预编译私有 IP CIDR 块，避免每次调用 isPrivateIP 时重复解析 -var privateNets []*net.IPNet - // CompiledIPRules 表示预编译的 IP 匹配规则。 // PatternCount 记录原始规则数量，用于保留“规则存在但全无效”时的行为语义。 type CompiledIPRules struct { @@ -75,23 +51,6 @@ type CompiledIPRules struct { PatternCount int } -func init() { - for _, cidr := range []string{ - "10.0.0.0/8", - "172.16.0.0/12", - "192.168.0.0/16", - "127.0.0.0/8", - "::1/128", - "fc00::/7", - } { - _, block, err := net.ParseCIDR(cidr) - if err != nil { - panic("invalid CIDR: " + cidr) - } - privateNets = append(privateNets, block) - } -} - // CompileIPRules 将 IP/CIDR 字符串规则预编译为可复用结构。 // 非法规则会被忽略，但 PatternCount 会保留原始规则条数。 func CompileIPRules(patterns []string) *CompiledIPRules { @@ -139,20 +98,6 @@ func matchesCompiledRules(parsedIP net.IP, rules *CompiledIPRules) bool { return false } -// isPrivateIP 检查 IP 是否为私有地址。 -func isPrivateIP(ipStr string) bool { - ip := net.ParseIP(ipStr) - if ip == nil { - return false - } - for _, block := range privateNets { - if block.Contains(ip) { - return true - } - } - return false -} - // MatchesPattern 检查 IP 是否匹配指定的模式（支持单个 IP 或 CIDR）。 // pattern 可以是： // - 单个 IP: "192.168.1.100" diff --git a/backend/internal/pkg/ip/ip_test.go b/backend/internal/pkg/ip/ip_test.go index 403b2d59e..4745c6f6e 100644 --- a/backend/internal/pkg/ip/ip_test.go +++ b/backend/internal/pkg/ip/ip_test.go @@ -10,56 +10,18 @@ import ( "github.com/stretchr/testify/require" ) -func TestIsPrivateIP(t *testing.T) { - tests := []struct { - name string - ip string - expected bool - }{ - // 私有 IPv4 - {"10.x 私有地址", "10.0.0.1", true}, - {"10.x 私有地址段末", "10.255.255.255", true}, - {"172.16.x 私有地址", "172.16.0.1", true}, - {"172.31.x 私有地址", "172.31.255.255", true}, - {"192.168.x 私有地址", "192.168.1.1", true}, - {"127.0.0.1 本地回环", "127.0.0.1", true}, - {"127.x 回环段", "127.255.255.255", true}, - - // 公网 IPv4 - {"8.8.8.8 公网 DNS", "8.8.8.8", false}, - {"1.1.1.1 公网", "1.1.1.1", false}, - {"172.15.255.255 非私有", "172.15.255.255", false}, - {"172.32.0.0 非私有", "172.32.0.0", false}, - {"11.0.0.1 公网", "11.0.0.1", false}, - - // IPv6 - {"::1 IPv6 回环", "::1", true}, - {"fc00:: IPv6 私有", "fc00::1", true}, - {"fd00:: IPv6 私有", "fd00::1", true}, - {"2001:db8::1 IPv6 公网", "2001:db8::1", false}, - - // 无效输入 - {"空字符串", "", false}, - {"非法字符串", "not-an-ip", false}, - {"不完整 IP", "192.168", false}, - } - - for _, tc := range tests { - t.Run(tc.name, func(t *testing.T) { - got := isPrivateIP(tc.ip) - require.Equal(t, tc.expected, got, "isPrivateIP(%q)", tc.ip) - }) - } -} - -func TestGetTrustedClientIPUsesGinClientIP(t *testing.T) { +func TestClientIPHelpersIgnoreSpoofedHeadersFromUntrustedPeer(t *testing.T) { gin.SetMode(gin.TestMode) r := gin.New() require.NoError(t, r.SetTrustedProxies(nil)) r.GET("/t", func(c *gin.Context) { - c.String(200, GetTrustedClientIP(c)) + c.JSON(200, gin.H{ + "client": GetClientIP(c), + "trusted": GetTrustedClientIP(c), + "security": GetSecurityClientIP(c), + }) }) w := httptest.NewRecorder() @@ -71,7 +33,7 @@ func TestGetTrustedClientIPUsesGinClientIP(t *testing.T) { r.ServeHTTP(w, req) require.Equal(t, 200, w.Code) - require.Equal(t, "9.9.9.9", w.Body.String()) + require.JSONEq(t, `{"client":"9.9.9.9","trusted":"9.9.9.9","security":"9.9.9.9"}`, w.Body.String()) } func TestCheckIPRestrictionWithCompiledRules(t *testing.T) { diff --git a/backend/internal/pkg/sysutil/restart.go b/backend/internal/pkg/sysutil/restart.go index 2146596fc..8925ecdb0 100644 --- a/backend/internal/pkg/sysutil/restart.go +++ b/backend/internal/pkg/sysutil/restart.go @@ -2,16 +2,63 @@ package sysutil import ( + "errors" + "fmt" "log" "os" "runtime" + "syscall" "time" ) +const restartSignalDelay = 100 * time.Millisecond + +type processSignaler interface { + Signal(os.Signal) error +} + +type processSignalScheduler func(time.Duration, func()) + +// scheduleProcessSignal defers a process signal long enough for the current +// HTTP response to be flushed. Signal failures are reported asynchronously +// because the scheduling caller has already returned by then. +func scheduleProcessSignal( + process processSignaler, + signal os.Signal, + delay time.Duration, + schedule processSignalScheduler, + onSignalError func(error), +) error { + if process == nil { + return errors.New("schedule process signal: nil process") + } + if signal == nil { + return errors.New("schedule process signal: nil signal") + } + if delay <= 0 { + return errors.New("schedule process signal: delay must be greater than zero") + } + if schedule == nil { + return errors.New("schedule process signal: nil scheduler") + } + if onSignalError == nil { + return errors.New("schedule process signal: nil signal error handler") + } + + schedule(delay, func() { + if err := process.Signal(signal); err != nil { + onSignalError(err) + } + }) + return nil +} + // RestartService triggers a service restart by gracefully exiting. // -// This relies on systemd's Restart=always configuration to automatically -// restart the service after it exits. This is the industry-standard approach: +// SIGHUP is handled by the main server lifecycle as a graceful restart request. +// After HTTP and application cleanup completes, the process exits non-zero so +// systemd Restart=on-failure (or Restart=always) starts the new process. +// This approach: // - Simple and reliable // - No sudo permissions needed // - No complex process management @@ -19,23 +66,29 @@ import ( // // Prerequisites: // - Linux OS with systemd -// - Service configured with Restart=always in systemd unit file +// - Service configured with Restart=on-failure or Restart=always func RestartService() error { if runtime.GOOS != "linux" { - log.Println("Service restart via exit only works on Linux with systemd") - return nil + return fmt.Errorf("service restart signal is only supported on Linux, current OS is %s", runtime.GOOS) } - log.Println("Initiating service restart by graceful exit...") - log.Println("systemd will automatically restart the service (Restart=always)") - - // Give a moment for logs to flush and response to be sent - go func() { - time.Sleep(100 * time.Millisecond) - os.Exit(0) - }() + process, err := os.FindProcess(os.Getpid()) + if err != nil { + return fmt.Errorf("find current process: %w", err) + } - return nil + log.Println("Scheduling graceful service restart...") + return scheduleProcessSignal( + process, + syscall.SIGHUP, + restartSignalDelay, + func(delay time.Duration, callback func()) { + time.AfterFunc(delay, callback) + }, + func(err error) { + log.Printf("Failed to signal graceful service restart: %v", err) + }, + ) } // RestartServiceAsync is a fire-and-forget version of RestartService. @@ -43,6 +96,6 @@ func RestartService() error { func RestartServiceAsync() { if err := RestartService(); err != nil { log.Printf("Service restart failed: %v", err) - log.Println("Please restart the service manually: sudo systemctl restart sub2api") + log.Println("Please restart the service manually through the configured process supervisor") } } diff --git a/backend/internal/pkg/sysutil/restart_test.go b/backend/internal/pkg/sysutil/restart_test.go new file mode 100644 index 000000000..1566538ab --- /dev/null +++ b/backend/internal/pkg/sysutil/restart_test.go @@ -0,0 +1,185 @@ +package sysutil + +import ( + "errors" + "os" + "runtime" + "strings" + "syscall" + "testing" + "time" +) + +type recordingProcessSignaler struct { + signals []os.Signal + signalErr error +} + +func (p *recordingProcessSignaler) Signal(signal os.Signal) error { + p.signals = append(p.signals, signal) + return p.signalErr +} + +func TestScheduleProcessSignalDefersSignal(t *testing.T) { + process := &recordingProcessSignaler{} + wantDelay := 175 * time.Millisecond + var ( + scheduleCalls int + scheduledDelay time.Duration + callback func() + handlerCalls int + ) + + err := scheduleProcessSignal( + process, + syscall.SIGHUP, + wantDelay, + func(delay time.Duration, scheduledCallback func()) { + scheduleCalls++ + scheduledDelay = delay + callback = scheduledCallback + }, + func(error) { + handlerCalls++ + }, + ) + if err != nil { + t.Fatalf("scheduleProcessSignal() error = %v", err) + } + if scheduleCalls != 1 { + t.Fatalf("scheduler calls = %d, want 1", scheduleCalls) + } + if scheduledDelay != wantDelay { + t.Fatalf("scheduled delay = %s, want %s", scheduledDelay, wantDelay) + } + if callback == nil { + t.Fatal("scheduler received a nil callback") + } + if len(process.signals) != 0 { + t.Fatalf("Signal called synchronously %d times, want 0", len(process.signals)) + } + if handlerCalls != 0 { + t.Fatalf("error handler called synchronously %d times, want 0", handlerCalls) + } + + callback() + + if len(process.signals) != 1 { + t.Fatalf("Signal calls after callback = %d, want 1", len(process.signals)) + } + if process.signals[0] != syscall.SIGHUP { + t.Fatalf("Signal argument = %v, want %v", process.signals[0], syscall.SIGHUP) + } + if handlerCalls != 0 { + t.Fatalf("error handler calls after successful signal = %d, want 0", handlerCalls) + } +} + +func TestScheduleProcessSignalReportsAsynchronousSignalError(t *testing.T) { + signalErr := errors.New("signal failed") + process := &recordingProcessSignaler{signalErr: signalErr} + var ( + callback func() + reportedError error + handlerCalls int + ) + + err := scheduleProcessSignal( + process, + syscall.SIGHUP, + time.Millisecond, + func(_ time.Duration, scheduledCallback func()) { + callback = scheduledCallback + }, + func(err error) { + handlerCalls++ + reportedError = err + }, + ) + if err != nil { + t.Fatalf("scheduleProcessSignal() error = %v", err) + } + if callback == nil { + t.Fatal("scheduler received a nil callback") + } + if handlerCalls != 0 { + t.Fatalf("error handler called before scheduled callback: %d times", handlerCalls) + } + + callback() + + if len(process.signals) != 1 { + t.Fatalf("Signal calls = %d, want 1", len(process.signals)) + } + if handlerCalls != 1 { + t.Fatalf("error handler calls = %d, want 1", handlerCalls) + } + if !errors.Is(reportedError, signalErr) { + t.Fatalf("reported error = %v, want %v", reportedError, signalErr) + } +} + +func TestScheduleProcessSignalRejectsInvalidArguments(t *testing.T) { + tests := []struct { + name string + process processSignaler + signal os.Signal + delay time.Duration + nilScheduler bool + nilErrHandler bool + }{ + {name: "nil process", process: nil, signal: syscall.SIGHUP, delay: time.Millisecond}, + {name: "nil signal", process: &recordingProcessSignaler{}, signal: nil, delay: time.Millisecond}, + {name: "zero delay", process: &recordingProcessSignaler{}, signal: syscall.SIGHUP, delay: 0}, + {name: "negative delay", process: &recordingProcessSignaler{}, signal: syscall.SIGHUP, delay: -time.Millisecond}, + {name: "nil scheduler", process: &recordingProcessSignaler{}, signal: syscall.SIGHUP, delay: time.Millisecond, nilScheduler: true}, + {name: "nil error handler", process: &recordingProcessSignaler{}, signal: syscall.SIGHUP, delay: time.Millisecond, nilErrHandler: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + scheduleCalls := 0 + handlerCalls := 0 + var scheduler processSignalScheduler = func(time.Duration, func()) { + scheduleCalls++ + } + if tt.nilScheduler { + scheduler = nil + } + var errorHandler func(error) = func(error) { + handlerCalls++ + } + if tt.nilErrHandler { + errorHandler = nil + } + + err := scheduleProcessSignal(tt.process, tt.signal, tt.delay, scheduler, errorHandler) + if err == nil { + t.Fatal("scheduleProcessSignal() error = nil, want validation error") + } + if scheduleCalls != 0 { + t.Fatalf("scheduler calls = %d, want 0", scheduleCalls) + } + if handlerCalls != 0 { + t.Fatalf("error handler calls = %d, want 0", handlerCalls) + } + if process, ok := tt.process.(*recordingProcessSignaler); ok && len(process.signals) != 0 { + t.Fatalf("Signal calls = %d, want 0", len(process.signals)) + } + }) + } +} + +func TestRestartServiceReturnsErrorOutsideLinux(t *testing.T) { + if runtime.GOOS == "linux" { + t.Skip("calling RestartService on Linux would signal the real test process") + } + + err := RestartService() + if err == nil { + t.Fatalf("RestartService() error = nil on %s, want unsupported-platform error", runtime.GOOS) + } + if !strings.Contains(err.Error(), runtime.GOOS) { + t.Fatalf("RestartService() error = %q, want current OS %q", err, runtime.GOOS) + } +} diff --git a/backend/internal/pkg/tlsfingerprint/dialer.go b/backend/internal/pkg/tlsfingerprint/dialer.go index c8d8369ff..5ba556c80 100644 --- a/backend/internal/pkg/tlsfingerprint/dialer.go +++ b/backend/internal/pkg/tlsfingerprint/dialer.go @@ -11,11 +11,90 @@ import ( "net" "net/http" "net/url" + "sync" + "time" utls "github.com/refraction-networking/utls" "golang.org/x/net/proxy" ) +// defaultDialStageTimeout bounds one network establishment stage when the +// caller does not provide a shorter deadline. It deliberately applies to TCP, +// proxy CONNECT/SOCKS negotiation, and the TLS handshake independently: a +// healthy connection follows the same protocol, while a broken proxy cannot +// hold an upstream request forever. +const defaultDialStageTimeout = 10 * time.Second + +func normalizeDialContext(ctx context.Context) context.Context { + if ctx == nil { + return context.Background() + } + return ctx +} + +func withDialStageTimeout(ctx context.Context, timeout time.Duration) (context.Context, context.CancelFunc) { + ctx = normalizeDialContext(ctx) + if timeout <= 0 { + timeout = defaultDialStageTimeout + } + return context.WithTimeout(ctx, timeout) +} + +// setConnDeadlineFromContext gives blocking proxy I/O the same deadline as +// its stage context. The fallback is useful for contexts without a deadline +// and makes the helper safe for custom connections used by callers/tests. +func setConnDeadlineFromContext(conn net.Conn, ctx context.Context, fallback time.Duration) error { + if conn == nil { + return fmt.Errorf("connection is nil") + } + deadline := time.Time{} + if fallback > 0 { + deadline = time.Now().Add(fallback) + } + if ctx != nil { + if ctxDeadline, ok := ctx.Deadline(); ok && (deadline.IsZero() || ctxDeadline.Before(deadline)) { + deadline = ctxDeadline + } + } + if deadline.IsZero() { + return nil + } + return conn.SetDeadline(deadline) +} + +func clearConnDeadline(conn net.Conn) { + if conn != nil { + _ = conn.SetDeadline(time.Time{}) + } +} + +// closeConnOnContextDone closes a connection while a manual proxy handshake +// is blocked in Read/Write. net.Conn operations do not all observe context +// cancellation themselves, so relying on context alone could leave a bad +// proxy holding the request goroutine. +func closeConnOnContextDone(ctx context.Context, conn net.Conn) func() { + if ctx == nil || conn == nil { + return func() {} + } + callbackDone := make(chan struct{}) + stop := context.AfterFunc(ctx, func() { + defer close(callbackDone) + _ = conn.Close() + }) + var cleanupOnce sync.Once + return func() { + cleanupOnce.Do(func() { + // context.AfterFunc's stop function does not wait when the + // callback has already started. Wait here so a successful dial + // cannot return a connection that is closed immediately after + // this cleanup function returns. + if !stop() { + <-callbackDone + } + }) + } +} + // Profile contains TLS fingerprint configuration. // All slice fields use built-in defaults when empty. type Profile struct { @@ -34,22 +113,25 @@ type Profile struct { // Dialer creates TLS connections with custom fingerprints. type Dialer struct { - profile *Profile - baseDialer func(ctx context.Context, network, addr string) (net.Conn, error) + profile *Profile + baseDialer func(ctx context.Context, network, addr string) (net.Conn, error) + stageTimeout time.Duration } // HTTPProxyDialer creates TLS connections through HTTP/HTTPS proxies with custom fingerprints. // It handles the CONNECT tunnel establishment before performing TLS handshake. type HTTPProxyDialer struct { - profile *Profile - proxyURL *url.URL + profile *Profile + proxyURL *url.URL + stageTimeout time.Duration } // SOCKS5ProxyDialer creates TLS connections through SOCKS5 proxies with custom fingerprints. // It uses golang.org/x/net/proxy to establish the SOCKS5 tunnel. type SOCKS5ProxyDialer struct { - profile *Profile - proxyURL *url.URL + profile *Profile + proxyURL *url.URL + stageTimeout time.Duration } // Default TLS fingerprint values captured from Claude Code (Node.js 24.x) @@ -121,26 +203,59 @@ var ( // If baseDialer is nil, direct TCP dial is used. func NewDialer(profile *Profile, baseDialer func(ctx context.Context, network, addr string) (net.Conn, error)) *Dialer { if baseDialer == nil { - baseDialer = (&net.Dialer{}).DialContext + baseDialer = (&net.Dialer{Timeout: defaultDialStageTimeout}).DialContext } - return &Dialer{profile: profile, baseDialer: baseDialer} + return &Dialer{profile: profile, baseDialer: baseDialer, stageTimeout: defaultDialStageTimeout} } // NewHTTPProxyDialer creates a new TLS fingerprint dialer that works through HTTP/HTTPS proxies. // It establishes a CONNECT tunnel before performing TLS handshake with custom fingerprint. func NewHTTPProxyDialer(profile *Profile, proxyURL *url.URL) *HTTPProxyDialer { - return &HTTPProxyDialer{profile: profile, proxyURL: proxyURL} + return &HTTPProxyDialer{profile: profile, proxyURL: proxyURL, stageTimeout: defaultDialStageTimeout} } // NewSOCKS5ProxyDialer creates a new TLS fingerprint dialer that works through SOCKS5 proxies. // It establishes a SOCKS5 tunnel before performing TLS handshake with custom fingerprint. func NewSOCKS5ProxyDialer(profile *Profile, proxyURL *url.URL) *SOCKS5ProxyDialer { - return &SOCKS5ProxyDialer{profile: profile, proxyURL: proxyURL} + return &SOCKS5ProxyDialer{profile: profile, proxyURL: proxyURL, stageTimeout: defaultDialStageTimeout} +} + +func dialStageTimeout(timeout time.Duration) time.Duration { + if timeout <= 0 { + return defaultDialStageTimeout + } + return timeout +} + +func (d *Dialer) timeout() time.Duration { + if d == nil { + return defaultDialStageTimeout + } + return dialStageTimeout(d.stageTimeout) +} + +func (d *HTTPProxyDialer) timeout() time.Duration { + if d == nil { + return defaultDialStageTimeout + } + return dialStageTimeout(d.stageTimeout) +} + +func (d *SOCKS5ProxyDialer) timeout() time.Duration { + if d == nil { + return defaultDialStageTimeout + } + return dialStageTimeout(d.stageTimeout) } // DialTLSContext establishes a TLS connection through SOCKS5 proxy with the configured fingerprint. // Flow: SOCKS5 CONNECT to target -> TLS handshake with utls on the tunnel func (d *SOCKS5ProxyDialer) DialTLSContext(ctx context.Context, network, addr string) (net.Conn, error) { + if d == nil || d.proxyURL == nil { + return nil, fmt.Errorf("SOCKS5 proxy dialer is not configured") + } + ctx = normalizeDialContext(ctx) + timeout := d.timeout() slog.Debug("tls_fingerprint_socks5_connecting", "proxy", d.proxyURL.Host, "target", addr) // Step 1: Create SOCKS5 dialer @@ -160,7 +275,10 @@ func (d *SOCKS5ProxyDialer) DialTLSContext(ctx context.Context, network, addr st proxyAddr = net.JoinHostPort(d.proxyURL.Hostname(), "1080") // Default SOCKS5 port } - socksDialer, err := proxy.SOCKS5("tcp", proxyAddr, auth, proxy.Direct) + // Pass a timeout-aware forwarding dialer so the proxy TCP connection is + // bounded as well as the SOCKS negotiation itself. + forwardDialer := &net.Dialer{Timeout: timeout} + socksDialer, err := proxy.SOCKS5("tcp", proxyAddr, auth, forwardDialer) if err != nil { slog.Debug("tls_fingerprint_socks5_dialer_failed", "error", err) return nil, fmt.Errorf("create SOCKS5 dialer: %w", err) @@ -168,7 +286,16 @@ func (d *SOCKS5ProxyDialer) DialTLSContext(ctx context.Context, network, addr st // Step 2: Establish SOCKS5 tunnel to target slog.Debug("tls_fingerprint_socks5_establishing_tunnel", "target", addr) - conn, err := socksDialer.Dial("tcp", addr) + contextDialer, ok := socksDialer.(proxy.ContextDialer) + if !ok { + return nil, fmt.Errorf("SOCKS5 dialer does not support context cancellation") + } + dialCtx, cancel := withDialStageTimeout(ctx, timeout) + // Preserve the historical SOCKS5 transport network ("tcp"); the caller's + // network value is still accepted by the public DialTLSContext contract, + // but was never used by the old implementation. + conn, err := contextDialer.DialContext(dialCtx, "tcp", addr) + cancel() if err != nil { slog.Debug("tls_fingerprint_socks5_connect_failed", "error", err) return nil, fmt.Errorf("SOCKS5 connect: %w", err) @@ -176,12 +303,17 @@ func (d *SOCKS5ProxyDialer) DialTLSContext(ctx context.Context, network, addr st slog.Debug("tls_fingerprint_socks5_tunnel_established") // Step 3: Perform TLS handshake on the tunnel with utls fingerprint - return performTLSHandshake(ctx, conn, d.profile, addr) + return performTLSHandshakeWithTimeout(ctx, conn, d.profile, addr, timeout) } // DialTLSContext establishes a TLS connection through HTTP proxy with the configured fingerprint. // Flow: TCP connect to proxy -> CONNECT tunnel -> TLS handshake with utls func (d *HTTPProxyDialer) DialTLSContext(ctx context.Context, network, addr string) (net.Conn, error) { + if d == nil || d.proxyURL == nil { + return nil, fmt.Errorf("HTTP proxy dialer is not configured") + } + ctx = normalizeDialContext(ctx) + timeout := d.timeout() slog.Debug("tls_fingerprint_http_proxy_connecting", "proxy", d.proxyURL.Host, "target", addr) // Step 1: TCP connect to proxy server @@ -197,14 +329,35 @@ func (d *HTTPProxyDialer) DialTLSContext(ctx context.Context, network, addr stri } } - dialer := &net.Dialer{} - conn, err := dialer.DialContext(ctx, "tcp", proxyAddr) + dialCtx, dialCancel := withDialStageTimeout(ctx, timeout) + dialer := &net.Dialer{Timeout: timeout} + conn, err := dialer.DialContext(dialCtx, "tcp", proxyAddr) + dialCancel() if err != nil { slog.Debug("tls_fingerprint_http_proxy_connect_failed", "error", err) return nil, fmt.Errorf("connect to proxy: %w", err) } slog.Debug("tls_fingerprint_http_proxy_connected", "proxy_addr", proxyAddr) + // The CONNECT exchange is performed manually (to preserve the custom TLS + // fingerprint), so explicitly bind both a socket deadline and context + // cancellation to the connection. + connectCtx, connectCancel := withDialStageTimeout(ctx, timeout) + stopContextClose := closeConnOnContextDone(connectCtx, conn) + if err := setConnDeadlineFromContext(conn, connectCtx, timeout); err != nil { + stopContextClose() + connectCancel() + _ = conn.Close() + return nil, fmt.Errorf("set proxy CONNECT deadline: %w", err) + } + finishConnect := func() error { + stopContextClose() + stageErr := connectCtx.Err() + connectCancel() + clearConnDeadline(conn) + return stageErr + } + // Step 2: Send CONNECT request to establish tunnel req := &http.Request{ Method: "CONNECT", @@ -223,8 +376,12 @@ func (d *HTTPProxyDialer) DialTLSContext(ctx context.Context, network, addr stri slog.Debug("tls_fingerprint_http_proxy_sending_connect", "target", addr) if err := req.Write(conn); err != nil { + stageErr := finishConnect() _ = conn.Close() slog.Debug("tls_fingerprint_http_proxy_write_failed", "error", err) + if stageErr != nil { + return nil, fmt.Errorf("proxy CONNECT stage: %w", stageErr) + } return nil, fmt.Errorf("write CONNECT request: %w", err) } @@ -232,30 +389,49 @@ func (d *HTTPProxyDialer) DialTLSContext(ctx context.Context, network, addr stri br := bufio.NewReader(conn) resp, err := http.ReadResponse(br, req) if err != nil { + stageErr := finishConnect() _ = conn.Close() slog.Debug("tls_fingerprint_http_proxy_read_response_failed", "error", err) + if stageErr != nil { + return nil, fmt.Errorf("proxy CONNECT stage: %w", stageErr) + } return nil, fmt.Errorf("read CONNECT response: %w", err) } // CONNECT response has no body; do not defer resp.Body.Close() as it wraps the // same conn that will be used for the TLS handshake. if resp.StatusCode != http.StatusOK { + stageErr := finishConnect() _ = conn.Close() slog.Debug("tls_fingerprint_http_proxy_connect_failed_status", "status_code", resp.StatusCode, "status", resp.Status) + if stageErr != nil { + return nil, fmt.Errorf("proxy CONNECT stage: %w", stageErr) + } return nil, fmt.Errorf("proxy CONNECT failed: %s", resp.Status) } slog.Debug("tls_fingerprint_http_proxy_tunnel_established") + if stageErr := finishConnect(); stageErr != nil { + _ = conn.Close() + return nil, fmt.Errorf("proxy CONNECT stage: %w", stageErr) + } // Step 4: Perform TLS handshake on the tunnel with utls fingerprint - return performTLSHandshake(ctx, conn, d.profile, addr) + return performTLSHandshakeWithTimeout(ctx, conn, d.profile, addr, timeout) } // DialTLSContext establishes a TLS connection with the configured fingerprint. // This method is designed to be used as http.Transport.DialTLSContext. func (d *Dialer) DialTLSContext(ctx context.Context, network, addr string) (net.Conn, error) { + if d == nil || d.baseDialer == nil { + return nil, fmt.Errorf("TLS dialer is not configured") + } + ctx = normalizeDialContext(ctx) + timeout := d.timeout() // Establish TCP connection using base dialer (supports proxy) slog.Debug("tls_fingerprint_dialing_tcp", "addr", addr) - conn, err := d.baseDialer(ctx, network, addr) + dialCtx, dialCancel := withDialStageTimeout(ctx, timeout) + conn, err := d.baseDialer(dialCtx, network, addr) + dialCancel() if err != nil { slog.Debug("tls_fingerprint_tcp_dial_failed", "error", err) return nil, err @@ -263,13 +439,30 @@ func (d *Dialer) DialTLSContext(ctx context.Context, network, addr string) (net. slog.Debug("tls_fingerprint_tcp_connected", "addr", addr) // Perform TLS handshake with utls fingerprint - return performTLSHandshake(ctx, conn, d.profile, addr) + return performTLSHandshakeWithTimeout(ctx, conn, d.profile, addr, timeout) } // performTLSHandshake performs the uTLS handshake on an established connection. // It builds a ClientHello spec from the profile, applies it, and completes the handshake. // On failure, conn is closed and an error is returned. func performTLSHandshake(ctx context.Context, conn net.Conn, profile *Profile, addr string) (net.Conn, error) { + return performTLSHandshakeWithTimeout(ctx, conn, profile, addr, defaultDialStageTimeout) +} + +func performTLSHandshakeWithTimeout(ctx context.Context, conn net.Conn, profile *Profile, addr string, timeout time.Duration) (net.Conn, error) { + if conn == nil { + return nil, fmt.Errorf("TLS handshake connection is nil") + } + ctx = normalizeDialContext(ctx) + handshakeCtx, handshakeCancel := withDialStageTimeout(ctx, timeout) + defer handshakeCancel() + stopContextClose := closeConnOnContextDone(handshakeCtx, conn) + defer stopContextClose() + if err := setConnDeadlineFromContext(conn, handshakeCtx, timeout); err != nil { + _ = conn.Close() + return nil, fmt.Errorf("set TLS handshake deadline: %w", err) + } + host, _, err := net.SplitHostPort(addr) if err != nil { host = addr @@ -283,10 +476,19 @@ func performTLSHandshake(ctx context.Context, conn net.Conn, profile *Profile, a return nil, fmt.Errorf("apply TLS preset: %w", err) } - if err := tlsConn.HandshakeContext(ctx); err != nil { + if err := tlsConn.HandshakeContext(handshakeCtx); err != nil { _ = conn.Close() return nil, fmt.Errorf("TLS handshake failed: %w", err) } + // Stop and, when necessary, join the cancellation callback before + // inspecting the stage context. Without this synchronization a timeout + // callback can close the just-completed connection after it is returned. + stopContextClose() + if err := handshakeCtx.Err(); err != nil { + _ = conn.Close() + return nil, fmt.Errorf("TLS handshake deadline exceeded: %w", err) + } + clearConnDeadline(tlsConn) state := tlsConn.ConnectionState() slog.Debug("tls_fingerprint_handshake_success", diff --git a/backend/internal/pkg/tlsfingerprint/dialer_test.go b/backend/internal/pkg/tlsfingerprint/dialer_test.go index 048418c94..0278667e6 100644 --- a/backend/internal/pkg/tlsfingerprint/dialer_test.go +++ b/backend/internal/pkg/tlsfingerprint/dialer_test.go @@ -14,10 +14,12 @@ import ( "context" "encoding/json" "io" + "net" "net/http" "net/url" "os" "strings" + "sync/atomic" "testing" "time" ) @@ -234,6 +236,228 @@ func TestSOCKS5ProxyDialerBasic(t *testing.T) { } } +func TestDialerBoundsCustomTCPDialStage(t *testing.T) { + dialer := NewDialer(&Profile{}, func(ctx context.Context, _, _ string) (net.Conn, error) { + <-ctx.Done() + return nil, ctx.Err() + }) + dialer.stageTimeout = 20 * time.Millisecond + + started := time.Now() + _, err := dialer.DialTLSContext(context.Background(), "tcp", "example.com:443") + if err == nil { + t.Fatal("expected bounded TCP dial to fail") + } + if elapsed := time.Since(started); elapsed > 500*time.Millisecond { + t.Fatalf("TCP dial exceeded stage bound: %s", elapsed) + } +} + +func TestTLSHandshakeBoundsBlockedPeer(t *testing.T) { + client, server := net.Pipe() + defer func() { + _ = client.Close() + _ = server.Close() + }() + + started := time.Now() + _, err := performTLSHandshakeWithTimeout(context.Background(), client, &Profile{}, "example.com:443", 20*time.Millisecond) + if err == nil { + t.Fatal("expected TLS handshake timeout") + } + if elapsed := time.Since(started); elapsed > 500*time.Millisecond { + t.Fatalf("TLS handshake exceeded stage bound: %s", elapsed) + } +} + +func TestHTTPProxyConnectHonorsStageTimeout(t *testing.T) { + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer listener.Close() + + accepted := make(chan net.Conn, 1) + go func() { + conn, acceptErr := listener.Accept() + if acceptErr == nil { + accepted <- conn + } + }() + + proxyURL, err := url.Parse("http://" + listener.Addr().String()) + if err != nil { + t.Fatal(err) + } + dialer := NewHTTPProxyDialer(&Profile{}, proxyURL) + dialer.stageTimeout = 20 * time.Millisecond + + started := time.Now() + _, err = dialer.DialTLSContext(context.Background(), "tcp", "example.com:443") + if err == nil { + t.Fatal("expected CONNECT timeout") + } + if elapsed := time.Since(started); elapsed > 500*time.Millisecond { + t.Fatalf("CONNECT exchange exceeded stage bound: %s", elapsed) + } + select { + case conn := <-accepted: + _ = conn.Close() + case <-time.After(500 * time.Millisecond): + t.Fatal("proxy did not accept the connection") + } +} + +func TestSOCKS5ConnectHonorsStageTimeout(t *testing.T) { + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer listener.Close() + + accepted := make(chan net.Conn, 1) + go func() { + conn, acceptErr := listener.Accept() + if acceptErr == nil { + accepted <- conn + } + }() + + proxyURL, err := url.Parse("socks5://" + listener.Addr().String()) + if err != nil { + t.Fatal(err) + } + dialer := NewSOCKS5ProxyDialer(&Profile{}, proxyURL) + dialer.stageTimeout = 20 * time.Millisecond + + started := time.Now() + _, err = dialer.DialTLSContext(context.Background(), "tcp", "example.com:443") + if err == nil { + t.Fatal("expected SOCKS5 negotiation timeout") + } + if elapsed := time.Since(started); elapsed > 500*time.Millisecond { + t.Fatalf("SOCKS5 negotiation exceeded stage bound: %s", elapsed) + } + select { + case conn := <-accepted: + _ = conn.Close() + case <-time.After(500 * time.Millisecond): + t.Fatal("SOCKS5 proxy did not accept the connection") + } +} + +func TestContextCancellationClosesBlockedConnection(t *testing.T) { + client, server := net.Pipe() + defer func() { + _ = client.Close() + _ = server.Close() + }() + + ctx, cancel := context.WithCancel(context.Background()) + stop := closeConnOnContextDone(ctx, client) + defer stop() + readDone := make(chan error, 1) + go func() { + buf := make([]byte, 1) + _, readErr := server.Read(buf) + readDone <- readErr + }() + + cancel() + select { + case readErr := <-readDone: + if readErr == nil { + t.Fatal("blocked connection unexpectedly returned without an error") + } + case <-time.After(500 * time.Millisecond): + t.Fatal("context cancellation did not close the connection") + } +} + +type trackingCloseConn struct { + net.Conn + closeCount atomic.Int32 +} + +func (c *trackingCloseConn) Close() error { + c.closeCount.Add(1) + return c.Conn.Close() +} + +type blockingCloseConn struct { + net.Conn + closeStarted chan struct{} + releaseClose chan struct{} +} + +func (c *blockingCloseConn) Close() error { + select { + case <-c.closeStarted: + default: + close(c.closeStarted) + } + <-c.releaseClose + return c.Conn.Close() +} + +func TestContextCleanupWaitsForStartedCloseCallback(t *testing.T) { + client, server := net.Pipe() + defer func() { + _ = client.Close() + _ = server.Close() + }() + conn := &blockingCloseConn{ + Conn: client, + closeStarted: make(chan struct{}), + releaseClose: make(chan struct{}), + } + ctx, cancel := context.WithCancel(context.Background()) + cleanup := closeConnOnContextDone(ctx, conn) + cancel() + + select { + case <-conn.closeStarted: + case <-time.After(500 * time.Millisecond): + t.Fatal("context cancellation callback did not start") + } + + cleanupDone := make(chan struct{}) + go func() { + cleanup() + close(cleanupDone) + }() + select { + case <-cleanupDone: + t.Fatal("cleanup returned before the started close callback completed") + case <-time.After(30 * time.Millisecond): + } + close(conn.releaseClose) + select { + case <-cleanupDone: + case <-time.After(500 * time.Millisecond): + t.Fatal("cleanup did not wait for the close callback to finish") + } +} + +func TestContextCleanupPreventsCloseAfterSuccessfulStage(t *testing.T) { + client, server := net.Pipe() + defer func() { + _ = client.Close() + _ = server.Close() + }() + conn := &trackingCloseConn{Conn: client} + ctx, cancel := context.WithCancel(context.Background()) + cleanup := closeConnOnContextDone(ctx, conn) + cleanup() + cancel() + // stop=true guarantees the callback is disassociated; the short wait also + // lets a wrongly implemented asynchronous callback surface deterministically. + time.Sleep(20 * time.Millisecond) + if got := conn.closeCount.Load(); got != 0 { + t.Fatalf("successful-stage cleanup was followed by an unexpected close: %d", got) + } +} + // TestBuildClientHelloSpec tests ClientHello spec construction. func TestBuildClientHelloSpec(t *testing.T) { // Test with nil profile (should use defaults) diff --git a/backend/internal/pkg/xai/billing.go b/backend/internal/pkg/xai/billing.go index 15b9c7e50..7dffdf29a 100644 --- a/backend/internal/pkg/xai/billing.go +++ b/backend/internal/pkg/xai/billing.go @@ -76,6 +76,8 @@ type BillingSummary struct { UsedPercent *float64 `json:"used_percent,omitempty"` Plan string `json:"plan,omitempty"` // SuperGrok | SuperGrok Heavy | "" StatusCode int `json:"status_code,omitempty"` + WeeklyStatusCode int `json:"weekly_status_code,omitempty"` + MonthlyStatusCode int `json:"monthly_status_code,omitempty"` Source string `json:"source,omitempty"` FetchedAt string `json:"fetched_at,omitempty"` UpdatedAt string `json:"updated_at,omitempty"` diff --git a/backend/internal/repository/account_repo.go b/backend/internal/repository/account_repo.go index 3874596f3..20bd78365 100644 --- a/backend/internal/repository/account_repo.go +++ b/backend/internal/repository/account_repo.go @@ -33,6 +33,7 @@ import ( "github.com/Wei-Shaw/sub2api/internal/service" "github.com/lib/pq" + "entgo.io/ent/dialect" entsql "entgo.io/ent/dialect/sql" "entgo.io/ent/dialect/sql/sqljson" ) @@ -64,6 +65,7 @@ var schedulerNeutralExtraKeyPrefixes = []string{ var schedulerNeutralExtraKeys = map[string]struct{}{ "codex_usage_updated_at": {}, + "grok_billing_snapshot": {}, "session_window_utilization": {}, } @@ -499,6 +501,65 @@ func (r *accountRepository) ExistsByCredentialField(ctx context.Context, key, va Exist(ctx) } +// GetOwnedOpenAIAgentIdentityByChatGPTAccountID returns the Agent Identity +// account owned by one user for one ChatGPT Team. The owner and auth-mode +// predicates are part of the database query so callers cannot observe or +// update another user's account through a cross-tenant lookup. A missing Team +// returns (nil, nil), which lets import callers distinguish absence from a +// persistence failure without treating the preflight lookup as an error. +func (r *accountRepository) GetOwnedOpenAIAgentIdentityByChatGPTAccountID( + ctx context.Context, + ownerUserID int64, + chatGPTAccountID string, +) (*service.Account, error) { + chatGPTAccountID = strings.TrimSpace(chatGPTAccountID) + if ownerUserID <= 0 || chatGPTAccountID == "" { + return nil, nil + } + trimFunction := "BTRIM" + if r.client.Driver().Dialect() != dialect.Postgres { + trimFunction = "TRIM" + } + + m, err := r.client.Account.Query(). + Where( + dbaccount.DeletedAtIsNil(), + dbaccount.OwnerUserIDEQ(ownerUserID), + dbaccount.PlatformEQ(service.PlatformOpenAI), + dbaccount.TypeEQ(service.AccountTypeOAuth), + func(selector *entsql.Selector) { + credentialsColumn := selector.C(dbaccount.FieldCredentials) + selector.Where(entsql.P(func(builder *entsql.Builder) { + builder.WriteString("LOWER(NULLIF("). + WriteString(trimFunction). + WriteString("("). + Ident(credentialsColumn). + WriteString("->>'auth_mode'), '')) = "). + Arg(strings.ToLower(service.OpenAIAuthModeAgentIdentity)). + WriteString(" AND NULLIF("). + WriteString(trimFunction). + WriteString("("). + Ident(credentialsColumn). + WriteString("->>'chatgpt_account_id'), '') = "). + Arg(chatGPTAccountID) + })) + }, + ). + Only(ctx) + if err != nil { + if dbent.IsNotFound(err) { + return nil, nil + } + return nil, err + } + + account := accountEntityToService(m) + if account == nil { + return nil, nil + } + return account, nil +} + func (r *accountRepository) IsAccountShareModeListingAccount(ctx context.Context, id int64) (bool, error) { if id <= 0 { return false, nil @@ -2064,24 +2125,19 @@ func (r *accountRepository) BindGroups(ctx context.Context, accountID int64, gro return err } - if len(groupIDs) == 0 { - if tx != nil { - return tx.Commit() + if len(groupIDs) > 0 { + builders := make([]*dbent.AccountGroupCreate, 0, len(groupIDs)) + for i, groupID := range groupIDs { + builders = append(builders, txClient.AccountGroup.Create(). + SetAccountID(accountID). + SetGroupID(groupID). + SetPriority(i+1), + ) } - return nil - } - builders := make([]*dbent.AccountGroupCreate, 0, len(groupIDs)) - for i, groupID := range groupIDs { - builders = append(builders, txClient.AccountGroup.Create(). - SetAccountID(accountID). - SetGroupID(groupID). - SetPriority(i+1), - ) - } - - if _, err := txClient.AccountGroup.CreateBulk(builders...).Save(ctx); err != nil { - return err + if _, err := txClient.AccountGroup.CreateBulk(builders...).Save(ctx); err != nil { + return err + } } if tx != nil { @@ -2889,6 +2945,12 @@ func (r *accountRepository) queryAccountsByGroup(ctx context.Context, groupID in } return nil, err } + if service.NormalizeGroupScope(group.Scope) == service.GroupScopePublic { + // Public groups may contain stale bindings while a share transition is + // being repaired. System-owned accounts keep their historical behavior, + // but user-owned accounts are schedulable publicly only after approval. + preds = append(preds, publicGroupSchedulableAccountPredicate()) + } requiredLevel := service.NormalizeRequiredAccountLevel(group.RequiredAccountLevel) if group.Platform == service.PlatformOpenAI && requiredLevel != "" { allowedLevels := service.OpenAISharedPoolAllowedAccountLevels(requiredLevel) @@ -2953,6 +3015,16 @@ func (r *accountRepository) queryAccountsByGroup(ctx context.Context, groupID in return r.accountsToService(ctx, accounts) } +func publicGroupSchedulableAccountPredicate() dbpredicate.Account { + return dbaccount.Or( + dbaccount.OwnerUserIDIsNil(), + dbaccount.And( + dbaccount.ShareModeEQ(service.AccountShareModePublic), + dbaccount.ShareStatusEQ(service.AccountShareStatusApproved), + ), + ) +} + func (r *accountRepository) accountsToService(ctx context.Context, accounts []*dbent.Account) ([]service.Account, error) { if len(accounts) == 0 { return []service.Account{}, nil diff --git a/backend/internal/repository/account_repo_agent_identity_test.go b/backend/internal/repository/account_repo_agent_identity_test.go new file mode 100644 index 000000000..3e06d2b61 --- /dev/null +++ b/backend/internal/repository/account_repo_agent_identity_test.go @@ -0,0 +1,134 @@ +package repository + +import ( + "context" + "database/sql" + "testing" + "time" + + dbent "github.com/Wei-Shaw/sub2api/ent" + "github.com/Wei-Shaw/sub2api/ent/enttest" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/stretchr/testify/require" + + "entgo.io/ent/dialect" + entsql "entgo.io/ent/dialect/sql" + _ "modernc.org/sqlite" +) + +func TestAccountRepositoryGetOwnedOpenAIAgentIdentityByChatGPTAccountID(t *testing.T) { + db, err := sql.Open("sqlite", "file:account_repo_agent_identity?mode=memory&cache=shared&_fk=1") + require.NoError(t, err) + t.Cleanup(func() { _ = db.Close() }) + db.SetMaxOpenConns(1) + _, err = db.Exec("PRAGMA foreign_keys = ON") + require.NoError(t, err) + + driver := entsql.OpenDB(dialect.SQLite, db) + client := enttest.NewClient(t, enttest.WithOptions(dbent.Driver(driver))) + t.Cleanup(func() { _ = client.Close() }) + repo := newAccountRepositoryWithSQL(client, db, nil) + ctx := context.Background() + + ownerA := createAgentIdentityRepositoryTestUser(t, ctx, client, "agent-owner-a@example.com") + ownerB := createAgentIdentityRepositoryTestUser(t, ctx, client, "agent-owner-b@example.com") + target := createAgentIdentityRepositoryTestAccount(t, ctx, client, ownerA, "target", map[string]any{ + "auth_mode": service.OpenAIAuthModeAgentIdentity, + "chatgpt_account_id": "team-a", + "chatgpt_user_id": "member-a", + "agent_runtime_id": "runtime-a", + }) + otherOwner := createAgentIdentityRepositoryTestAccount(t, ctx, client, ownerB, "other-owner", map[string]any{ + "auth_mode": service.OpenAIAuthModeAgentIdentity, + "chatgpt_account_id": "team-a", + "chatgpt_user_id": "member-b", + "agent_runtime_id": "runtime-b", + }) + createAgentIdentityRepositoryTestAccount(t, ctx, client, ownerA, "normal-oauth", map[string]any{ + "access_token": "token", + "chatgpt_account_id": "team-a", + }) + createAgentIdentityRepositoryTestAccount(t, ctx, client, ownerA, "other-team", map[string]any{ + "auth_mode": service.OpenAIAuthModeAgentIdentity, + "chatgpt_account_id": "team-b", + "chatgpt_user_id": "member-a", + "agent_runtime_id": "runtime-c", + }) + nonCanonical := createAgentIdentityRepositoryTestAccount(t, ctx, client, ownerA, "non-canonical", map[string]any{ + "auth_mode": " AGENTIDENTITY ", + "chatgpt_account_id": " team-non-canonical ", + "chatgpt_user_id": "member-non-canonical", + "agent_runtime_id": "runtime-non-canonical", + }) + deleted := createAgentIdentityRepositoryTestAccount(t, ctx, client, ownerA, "deleted", map[string]any{ + "auth_mode": service.OpenAIAuthModeAgentIdentity, + "chatgpt_account_id": "team-deleted", + "chatgpt_user_id": "member-a", + "agent_runtime_id": "runtime-deleted", + }) + _, err = client.Account.UpdateOneID(deleted.ID).SetDeletedAt(time.Now().UTC()).Save(ctx) + require.NoError(t, err) + + got, err := repo.GetOwnedOpenAIAgentIdentityByChatGPTAccountID(ctx, ownerA, " team-a ") + require.NoError(t, err) + require.Equal(t, target.ID, got.ID) + require.Equal(t, int64(ownerA), *got.OwnerUserID) + require.Equal(t, "member-a", got.GetCredential("chatgpt_user_id")) + + got, err = repo.GetOwnedOpenAIAgentIdentityByChatGPTAccountID(ctx, ownerB, "team-a") + require.NoError(t, err) + require.Equal(t, otherOwner.ID, got.ID) + + got, err = repo.GetOwnedOpenAIAgentIdentityByChatGPTAccountID(ctx, ownerA, "team-non-canonical") + require.NoError(t, err) + require.Equal(t, nonCanonical.ID, got.ID) + + for _, test := range []struct { + name string + ownerID int64 + accountID string + }{ + {name: "unknown team", ownerID: ownerA, accountID: "team-missing"}, + {name: "soft deleted", ownerID: ownerA, accountID: "team-deleted"}, + {name: "invalid owner", ownerID: 0, accountID: "team-a"}, + {name: "empty team", ownerID: ownerA, accountID: " "}, + } { + t.Run(test.name, func(t *testing.T) { + account, err := repo.GetOwnedOpenAIAgentIdentityByChatGPTAccountID(ctx, test.ownerID, test.accountID) + require.Nil(t, account) + require.NoError(t, err) + }) + } +} + +func createAgentIdentityRepositoryTestUser(t *testing.T, ctx context.Context, client *dbent.Client, email string) int64 { + t.Helper() + user, err := client.User.Create(). + SetEmail(email). + SetPasswordHash("test-password-hash"). + SetRole(service.RoleUser). + SetStatus(service.StatusActive). + Save(ctx) + require.NoError(t, err) + return user.ID +} + +func createAgentIdentityRepositoryTestAccount( + t *testing.T, + ctx context.Context, + client *dbent.Client, + ownerUserID int64, + name string, + credentials map[string]any, +) *dbent.Account { + t.Helper() + account, err := client.Account.Create(). + SetName(name). + SetPlatform(service.PlatformOpenAI). + SetType(service.AccountTypeOAuth). + SetOwnerUserID(ownerUserID). + SetCredentials(credentials). + Save(ctx) + require.NoError(t, err) + return account +} diff --git a/backend/internal/repository/account_repo_grok_managed_extra_test.go b/backend/internal/repository/account_repo_grok_managed_extra_test.go new file mode 100644 index 000000000..10da30dcd --- /dev/null +++ b/backend/internal/repository/account_repo_grok_managed_extra_test.go @@ -0,0 +1,21 @@ +//go:build unit + +package repository + +import ( + "testing" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/stretchr/testify/require" +) + +func TestGrokBillingSnapshotIsSchedulerNeutral(t *testing.T) { + t.Parallel() + + require.False(t, shouldEnqueueSchedulerOutboxForExtraUpdates(map[string]any{ + "grok_billing_snapshot": map[string]any{"status_code": float64(200)}, + })) + require.True(t, shouldEnqueueSchedulerOutboxForExtraUpdates(map[string]any{ + service.GrokMediaEligibleExtraKey: true, + }), "operator eligibility overrides must still rebuild affected scheduler buckets") +} diff --git a/backend/internal/repository/account_repo_integration_test.go b/backend/internal/repository/account_repo_integration_test.go index be390be65..572bf5f11 100644 --- a/backend/internal/repository/account_repo_integration_test.go +++ b/backend/internal/repository/account_repo_integration_test.go @@ -632,6 +632,17 @@ func (s *AccountRepoSuite) TestBindGroups_EmptyList() { groups, err := s.repo.GetGroups(s.ctx, account.ID) s.Require().NoError(err) s.Require().Empty(groups, "expected 0 groups after binding empty list") + + var outboxCount int + err = scanSingleRow( + s.ctx, + s.repo.sql, + "SELECT COUNT(*) FROM scheduler_outbox WHERE event_type = $1 AND account_id = $2", + []any{service.SchedulerOutboxEventAccountGroupsChanged, account.ID}, + &outboxCount, + ) + s.Require().NoError(err) + s.Require().Equal(1, outboxCount, "clearing groups must invalidate scheduler buckets that contained the account") } // --- Schedulable --- diff --git a/backend/internal/repository/account_repo_public_share_scheduler_test.go b/backend/internal/repository/account_repo_public_share_scheduler_test.go new file mode 100644 index 000000000..262d6ef03 --- /dev/null +++ b/backend/internal/repository/account_repo_public_share_scheduler_test.go @@ -0,0 +1,79 @@ +package repository + +import ( + "context" + "database/sql" + "testing" + + dbent "github.com/Wei-Shaw/sub2api/ent" + "github.com/Wei-Shaw/sub2api/ent/enttest" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/stretchr/testify/require" + + "entgo.io/ent/dialect" + entsql "entgo.io/ent/dialect/sql" + _ "modernc.org/sqlite" +) + +func TestPublicGroupSchedulableAccountPredicateRequiresApprovedOwnedShare(t *testing.T) { + db, err := sql.Open("sqlite", "file:account_repo_public_share_predicate?mode=memory&cache=shared&_fk=1") + require.NoError(t, err) + t.Cleanup(func() { _ = db.Close() }) + db.SetMaxOpenConns(1) + _, err = db.Exec("PRAGMA foreign_keys = ON") + require.NoError(t, err) + + driver := entsql.OpenDB(dialect.SQLite, db) + client := enttest.NewClient(t, enttest.WithOptions(dbent.Driver(driver))) + t.Cleanup(func() { _ = client.Close() }) + ctx := context.Background() + + owner, err := client.User.Create(). + SetEmail("public-share-predicate@example.com"). + SetPasswordHash("test-password-hash"). + SetRole(service.RoleUser). + SetStatus(service.StatusActive). + Save(ctx) + require.NoError(t, err) + + ownerID := owner.ID + accounts := []*dbent.Account{ + createPublicSharePredicateTestAccount(t, ctx, client, "system", nil, service.AccountShareModePrivate, service.AccountShareStatusApproved), + createPublicSharePredicateTestAccount(t, ctx, client, "owned-approved", &ownerID, service.AccountShareModePublic, service.AccountShareStatusApproved), + createPublicSharePredicateTestAccount(t, ctx, client, "owned-pending", &ownerID, service.AccountShareModePublic, service.AccountShareStatusPending), + createPublicSharePredicateTestAccount(t, ctx, client, "owned-suspended", &ownerID, service.AccountShareModePublic, service.AccountShareStatusSuspended), + createPublicSharePredicateTestAccount(t, ctx, client, "owned-private", &ownerID, service.AccountShareModePrivate, service.AccountShareStatusApproved), + } + + ids, err := client.Account.Query(). + Where(publicGroupSchedulableAccountPredicate()). + IDs(ctx) + require.NoError(t, err) + require.ElementsMatch(t, []int64{accounts[0].ID, accounts[1].ID}, ids) +} + +func createPublicSharePredicateTestAccount( + t *testing.T, + ctx context.Context, + client *dbent.Client, + name string, + ownerUserID *int64, + shareMode string, + shareStatus string, +) *dbent.Account { + t.Helper() + builder := client.Account.Create(). + SetName(name). + SetPlatform(service.PlatformOpenAI). + SetType(service.AccountTypeOAuth). + SetShareMode(shareMode). + SetShareStatus(shareStatus). + SetStatus(service.StatusActive). + SetSchedulable(true) + if ownerUserID != nil { + builder.SetOwnerUserID(*ownerUserID) + } + account, err := builder.Save(ctx) + require.NoError(t, err) + return account +} diff --git a/backend/internal/repository/ent.go b/backend/internal/repository/ent.go index 51b2d3282..569e3ee0e 100644 --- a/backend/internal/repository/ent.go +++ b/backend/internal/repository/ent.go @@ -5,13 +5,13 @@ package repository import ( "context" "database/sql" + "errors" "fmt" "time" "github.com/Wei-Shaw/sub2api/ent" "github.com/Wei-Shaw/sub2api/internal/config" "github.com/Wei-Shaw/sub2api/internal/pkg/timezone" - "github.com/Wei-Shaw/sub2api/migrations" "entgo.io/ent/dialect" entsql "entgo.io/ent/dialect/sql" @@ -36,31 +36,10 @@ import ( // - *sql.DB: 底层的 SQL 数据库连接，可用于直接执行原生 SQL // - error: 初始化过程中的错误 func InitEnt(cfg *config.Config) (*ent.Client, *sql.DB, error) { - // 优先初始化时区设置，确保所有时间操作使用统一的时区。 - // 这对于跨时区部署和日志时间戳的一致性至关重要。 - if err := timezone.Init(cfg.Timezone); err != nil { + drv, err := openConfiguredPostgresDriver(cfg) + if err != nil { return nil, nil, err } - - // 构建包含时区信息的数据库连接字符串 (DSN)。 - // 时区信息会传递给 PostgreSQL，确保数据库层面的时间处理正确。 - dsn := cfg.Database.DSNWithTimezone(cfg.Timezone) - - // 仅在显式开启 Server-Timing 时包装 driver，默认路径保持零额外驱动开销。 - var drv *entsql.Driver - if cfg.Server.EnableServerTiming { - connector, err := pq.NewConnector(dsn) - if err != nil { - return nil, nil, fmt.Errorf("create PostgreSQL connector: %w", err) - } - drv = entsql.OpenDB(dialect.Postgres, sql.OpenDB(newServerTimingConnector(connector))) - } else { - var err error - drv, err = entsql.Open(dialect.Postgres, dsn) - if err != nil { - return nil, nil, err - } - } applyDBPoolSettings(drv.DB(), cfg) // 确保数据库 schema 已准备就绪。 @@ -68,7 +47,7 @@ func InitEnt(cfg *config.Config) (*ent.Client, *sql.DB, error) { // 这种方式比 Ent 的自动迁移更可控，支持复杂的迁移场景。 migrationCtx, cancel := context.WithTimeout(context.Background(), 10*time.Minute) defer cancel() - if err := applyMigrationsFS(migrationCtx, drv.DB(), migrations.FS); err != nil { + if err := ApplyMigrations(migrationCtx, drv.DB()); err != nil { _ = drv.Close() // 迁移失败时关闭驱动，避免资源泄露 return nil, nil, err } @@ -106,3 +85,65 @@ func InitEnt(cfg *config.Config) (*ent.Client, *sql.DB, error) { return client, drv.DB(), nil } + +// ApplyConfiguredMigrations 使用应用配置连接 PostgreSQL，仅执行内嵌数据库迁移并关闭连接。 +// 该入口不会初始化 Ent 客户端、Redis、HTTP 服务或后台任务。 +func ApplyConfiguredMigrations(ctx context.Context, cfg *config.Config) (err error) { + if ctx == nil { + return errors.New("nil migration context") + } + + drv, err := openConfiguredPostgresDriver(cfg) + if err != nil { + return err + } + // migrate-only 使用独占单连接池，确保会话级 advisory lock 与全部迁移语句 + // 落在同一 PostgreSQL session，并避免迁移过程中连接按生命周期被轮换。 + drv.DB().SetMaxOpenConns(1) + drv.DB().SetMaxIdleConns(1) + drv.DB().SetConnMaxLifetime(0) + drv.DB().SetConnMaxIdleTime(0) + defer func() { + if closeErr := drv.Close(); closeErr != nil && err == nil { + err = fmt.Errorf("close PostgreSQL migration connection: %w", closeErr) + } + }() + + if err := ApplyMigrations(ctx, drv.DB()); err != nil { + return fmt.Errorf("apply configured migrations: %w", err) + } + return nil +} + +func openConfiguredPostgresDriver(cfg *config.Config) (*entsql.Driver, error) { + if cfg == nil { + return nil, errors.New("nil config") + } + + // 优先初始化时区设置，确保所有时间操作使用统一的时区。 + // 这对于跨时区部署和日志时间戳的一致性至关重要。 + if err := timezone.Init(cfg.Timezone); err != nil { + return nil, err + } + + // 构建包含时区信息的数据库连接字符串 (DSN)。 + // 时区信息会传递给 PostgreSQL，确保数据库层面的时间处理正确。 + dsn := cfg.Database.DSNWithTimezone(cfg.Timezone) + + // 仅在显式开启 Server-Timing 时包装 driver，默认路径保持零额外驱动开销。 + var drv *entsql.Driver + if cfg.Server.EnableServerTiming { + connector, err := pq.NewConnector(dsn) + if err != nil { + return nil, fmt.Errorf("create PostgreSQL connector: %w", err) + } + drv = entsql.OpenDB(dialect.Postgres, sql.OpenDB(newServerTimingConnector(connector))) + } else { + var err error + drv, err = entsql.Open(dialect.Postgres, dsn) + if err != nil { + return nil, err + } + } + return drv, nil +} diff --git a/backend/internal/repository/ent_migrate_only_test.go b/backend/internal/repository/ent_migrate_only_test.go new file mode 100644 index 000000000..f106849f8 --- /dev/null +++ b/backend/internal/repository/ent_migrate_only_test.go @@ -0,0 +1,32 @@ +package repository + +import ( + "context" + "strings" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" +) + +func TestApplyConfiguredMigrationsRejectsNilInputs(t *testing.T) { + var nilContext context.Context + + tests := []struct { + name string + ctx context.Context + cfg *config.Config + wantErr string + }{ + {name: "nil context", ctx: nilContext, cfg: &config.Config{}, wantErr: "nil migration context"}, + {name: "nil config", ctx: context.Background(), wantErr: "nil config"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := ApplyConfiguredMigrations(tt.ctx, tt.cfg) + if err == nil || !strings.Contains(err.Error(), tt.wantErr) { + t.Fatalf("ApplyConfiguredMigrations() error = %v, want containing %q", err, tt.wantErr) + } + }) + } +} diff --git a/backend/internal/repository/github_release_service.go b/backend/internal/repository/github_release_service.go index ad1f22e39..2e1e77c3d 100644 --- a/backend/internal/repository/github_release_service.go +++ b/backend/internal/repository/github_release_service.go @@ -7,6 +7,7 @@ import ( "io" "log/slog" "net/http" + "net/url" "os" "strings" "time" @@ -18,6 +19,7 @@ import ( type githubReleaseClient struct { httpClient *http.Client downloadHTTPClient *http.Client + updateGitHubToken string } type githubReleaseClientError struct { @@ -43,6 +45,8 @@ func NewGitHubReleaseClient(proxyURL string, allowDirectOnProxyError bool) servi } sharedClient = &http.Client{Timeout: 30 * time.Second} } + apiClient := cloneGitHubHTTPClient(sharedClient) + apiClient.CheckRedirect = githubAPICheckRedirect(apiClient.CheckRedirect) // 下载客户端需要更长的超时时间 downloadClient, err := httpclient.GetClient(httpclient.Options{ @@ -56,13 +60,50 @@ func NewGitHubReleaseClient(proxyURL string, allowDirectOnProxyError bool) servi } downloadClient = &http.Client{Timeout: 10 * time.Minute} } + downloadClient = cloneGitHubHTTPClient(downloadClient) return &githubReleaseClient{ - httpClient: sharedClient, + httpClient: apiClient, downloadHTTPClient: downloadClient, + updateGitHubToken: strings.TrimSpace(os.Getenv("UPDATE_GITHUB_TOKEN")), } } +func cloneGitHubHTTPClient(client *http.Client) *http.Client { + cloned := *client + return &cloned +} + +func isGitHubAPIURL(parsed *url.URL) bool { + return parsed != nil && strings.EqualFold(parsed.Scheme, "https") && parsed.User == nil && + strings.EqualFold(parsed.Host, "api.github.com") +} + +func githubAPICheckRedirect(previous func(*http.Request, []*http.Request) error) func(*http.Request, []*http.Request) error { + return func(req *http.Request, via []*http.Request) error { + if !isGitHubAPIURL(req.URL) { + req.Header.Del("Authorization") + } + if previous != nil { + return previous(req, via) + } + return nil + } +} + +func (c *githubReleaseClient) newAPIRequest(ctx context.Context, requestURL string) (*http.Request, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, requestURL, nil) + if err != nil { + return nil, err + } + req.Header.Set("Accept", "application/vnd.github.v3+json") + req.Header.Set("User-Agent", "Sub2API-Updater") + if c.updateGitHubToken != "" && isGitHubAPIURL(req.URL) { + req.Header.Set("Authorization", "Bearer "+c.updateGitHubToken) + } + return req, nil +} + func (c *githubReleaseClientError) FetchLatestRelease(ctx context.Context, repo string) (*service.GitHubRelease, error) { return nil, c.err } @@ -78,12 +119,10 @@ func (c *githubReleaseClientError) FetchChecksumFile(ctx context.Context, url st func (c *githubReleaseClient) FetchLatestRelease(ctx context.Context, repo string) (*service.GitHubRelease, error) { url := fmt.Sprintf("https://api.github.com/repos/%s/releases/latest", repo) - req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + req, err := c.newAPIRequest(ctx, url) if err != nil { return nil, err } - req.Header.Set("Accept", "application/vnd.github.v3+json") - req.Header.Set("User-Agent", "Sub2API-Updater") resp, err := c.httpClient.Do(req) if err != nil { diff --git a/backend/internal/repository/github_release_service_test.go b/backend/internal/repository/github_release_service_test.go index d375a1930..a5277a046 100644 --- a/backend/internal/repository/github_release_service_test.go +++ b/backend/internal/repository/github_release_service_test.go @@ -44,6 +44,40 @@ func newTestGitHubReleaseClient() *githubReleaseClient { } } +func TestGitHubReleaseClientAPIRequestAuthorization(t *testing.T) { + tests := []struct { + name string + url string + wantAuth string + }{ + {name: "exact HTTPS authority", url: "https://api.github.com/repos/test/repo", wantAuth: "Bearer update-secret"}, + {name: "HTTP", url: "http://api.github.com/repos/test/repo"}, + {name: "subdomain", url: "https://sub.api.github.com/repos/test/repo"}, + {name: "userinfo", url: "https://user@api.github.com/repos/test/repo"}, + {name: "explicit port", url: "https://api.github.com:443/repos/test/repo"}, + {name: "different host", url: "https://github.com/test/repo"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + client := newTestGitHubReleaseClient() + client.updateGitHubToken = "update-secret" + req, err := client.newAPIRequest(context.Background(), tt.url) + require.NoError(t, err) + require.Equal(t, tt.wantAuth, req.Header.Get("Authorization")) + }) + } +} + +func TestGitHubReleaseClientRedirectStripsAuthorizationAcrossHosts(t *testing.T) { + req, err := http.NewRequest(http.MethodGet, "https://objects.githubusercontent.com/asset", nil) + require.NoError(t, err) + req.Header.Set("Authorization", "Bearer update-secret") + + require.NoError(t, githubAPICheckRedirect(nil)(req, nil)) + require.Empty(t, req.Header.Get("Authorization")) +} + func (s *GitHubReleaseServiceSuite) SetupTest() { s.tempDir = s.T().TempDir() } diff --git a/backend/internal/repository/http_upstream.go b/backend/internal/repository/http_upstream.go index 4363d762b..9c3f06fc1 100644 --- a/backend/internal/repository/http_upstream.go +++ b/backend/internal/repository/http_upstream.go @@ -52,6 +52,10 @@ const ( // defaultResponseHeaderTimeout: 默认等待响应头超时时间（5分钟） // LLM 请求可能排队较久，需要较长超时 defaultResponseHeaderTimeout = 300 * time.Second + // 冷连接阶段必须有界；正常热连接不会触发这些超时。 + defaultUpstreamDialTimeout = 10 * time.Second + defaultUpstreamTLSHandshakeTimeout = 10 * time.Second + defaultUpstreamExpectContinueTimeout = 1 * time.Second // defaultMaxUpstreamClients: 默认最大客户端缓存数量 // 超出后会淘汰最久未使用的客户端 defaultMaxUpstreamClients = 5000 @@ -170,7 +174,7 @@ func NewHTTPUpstream(cfg *config.Config) service.HTTPUpstream { // - error: 请求错误 // // 注意: -// - 调用方必须关闭 resp.Body，否则会导致 inFlight 计数泄漏 +// - 调用方仍应关闭 resp.Body；完整读取到终态时也会自动释放 inFlight // - inFlight > 0 的客户端不会被淘汰，确保活跃请求不被中断 func (s *httpUpstreamService) Do(req *http.Request, proxyURL string, accountID int64, accountConcurrency int) (*http.Response, error) { applyGrokCLIProxyHeaders(req) @@ -188,8 +192,9 @@ func (s *httpUpstreamService) Do(req *http.Request, proxyURL string, accountID i return nil, err } - // 执行请求 - resp, err := servertiming.Do(entry.client, req) + // 执行请求。对携带凭证且禁止跳转的请求，仅浅拷贝 client，避免修改共享连接池。 + client := httpClientForUpstreamRequest(entry.client, req) + resp, err := servertiming.Do(client, req) if err != nil { s.recordOpenAIHTTP2Failure(profile, entry.protocolMode, entry.proxyKey, err) // 请求失败，立即减少计数 @@ -202,7 +207,7 @@ func (s *httpUpstreamService) Do(req *http.Request, proxyURL string, accountID i // 如果上游返回了压缩内容，解压后再交给业务层 decompressResponseBody(resp) - // 包装响应体，在关闭时自动减少计数并更新时间戳 + // 包装响应体，在读取结束或关闭时自动减少计数并更新时间戳 // 这确保了流式响应（如 SSE）在完全读取前不会被淘汰 resp.Body = wrapTrackedBody(resp.Body, func() { atomic.AddInt64(&entry.inFlight, -1) @@ -246,7 +251,8 @@ func (s *httpUpstreamService) DoWithTLS(req *http.Request, proxyURL string, acco return nil, err } - resp, err := servertiming.Do(entry.client, req) + client := httpClientForUpstreamRequest(entry.client, req) + resp, err := servertiming.Do(client, req) if err != nil { atomic.AddInt64(&entry.inFlight, -1) atomic.StoreInt64(&entry.lastUsed, time.Now().UnixNano()) @@ -264,6 +270,17 @@ func (s *httpUpstreamService) DoWithTLS(req *http.Request, proxyURL string, acco return resp, nil } +func httpClientForUpstreamRequest(client *http.Client, req *http.Request) *http.Client { + if client == nil || req == nil || !service.HTTPUpstreamRedirectsDisabled(req.Context()) { + return client + } + clone := *client + clone.CheckRedirect = func(*http.Request, []*http.Request) error { + return http.ErrUseLastResponse + } + return &clone +} + // applyGrokCLIProxyHeaders applies the official Grok Build client identity at // the final shared transport boundary and leaves direct api.x.ai traffic unchanged. func applyGrokCLIProxyHeaders(req *http.Request) { @@ -1051,6 +1068,9 @@ func defaultPoolSettings(cfg *config.Config) poolSettings { // - ResponseHeaderTimeout: 等待响应头超时（不影响流式传输） func buildUpstreamTransport(settings poolSettings, proxyURL *url.URL, protocolMode string) (*http.Transport, error) { transport := &http.Transport{ + DialContext: (&net.Dialer{Timeout: defaultUpstreamDialTimeout, KeepAlive: 30 * time.Second}).DialContext, + TLSHandshakeTimeout: defaultUpstreamTLSHandshakeTimeout, + ExpectContinueTimeout: defaultUpstreamExpectContinueTimeout, MaxIdleConns: settings.maxIdleConns, MaxIdleConnsPerHost: settings.maxIdleConnsPerHost, MaxConnsPerHost: settings.maxConnsPerHost, @@ -1105,6 +1125,12 @@ func enableOpenAIHTTP2KeepAlive(transport *http.Transport) (*http2.Transport, er // - socks5: SOCKS5 代理，使用 SOCKS5ProxyDialer（SOCKS5 隧道 + utls 握手） func buildUpstreamTransportWithTLSFingerprint(settings poolSettings, proxyURL *url.URL, profile *tlsfingerprint.Profile) (*http.Transport, error) { transport := &http.Transport{ + // DialTLSContext below enforces the same bounds for the custom-fingerprint + // path; keep these fields populated for transport-level diagnostics and + // for any future fallback to the standard dialer. + DialContext: (&net.Dialer{Timeout: defaultUpstreamDialTimeout, KeepAlive: 30 * time.Second}).DialContext, + TLSHandshakeTimeout: defaultUpstreamTLSHandshakeTimeout, + ExpectContinueTimeout: defaultUpstreamExpectContinueTimeout, MaxIdleConns: settings.maxIdleConns, MaxIdleConnsPerHost: settings.maxIdleConnsPerHost, MaxConnsPerHost: settings.maxConnsPerHost, @@ -1145,38 +1171,55 @@ func buildUpstreamTransportWithTLSFingerprint(settings poolSettings, proxyURL *u return transport, nil } -// trackedBody 带跟踪功能的响应体包装器 -// 在 Close 时执行回调，用于更新请求计数 +// trackedBody 带跟踪功能的响应体包装器。 +// 响应读取结束或关闭时执行回调，用于更新请求计数。 type trackedBody struct { io.ReadCloser // 原始响应体 once sync.Once - onClose func() // 关闭时的回调函数 + onDone func() +} + +// Read 在响应体到达终态后立即释放占用。 +// 这避免调用方完整读取响应体、但遗漏 Close 时永久泄漏 inFlight。 +func (b *trackedBody) Read(p []byte) (int, error) { + n, err := b.ReadCloser.Read(p) + if err != nil { + b.done() + } + return n, err } -// Close 关闭响应体并执行回调 -// 使用 sync.Once 确保回调只执行一次 +// Close 关闭响应体并执行回调。 +// 使用 sync.Once 确保 Read 终态和 Close 最多释放一次。 func (b *trackedBody) Close() error { err := b.ReadCloser.Close() - if b.onClose != nil { - b.once.Do(b.onClose) - } + b.done() return err } -// wrapTrackedBody 包装响应体以跟踪关闭事件 -// 用于在响应体关闭时更新 inFlight 计数 +func (b *trackedBody) done() { + if b.onDone != nil { + b.once.Do(b.onDone) + } +} + +// wrapTrackedBody 包装响应体以跟踪完成事件。 +// 用于在响应体读取结束或关闭时更新 inFlight 计数。 // // 参数: // - body: 原始响应体 -// - onClose: 关闭时的回调函数 +// - onDone: 读取结束或关闭时的回调函数 // // 返回: // - io.ReadCloser: 包装后的响应体 -func wrapTrackedBody(body io.ReadCloser, onClose func()) io.ReadCloser { +func wrapTrackedBody(body io.ReadCloser, onDone func()) io.ReadCloser { if body == nil { + if onDone != nil { + onDone() + } return body } - return &trackedBody{ReadCloser: body, onClose: onClose} + return &trackedBody{ReadCloser: body, onDone: onDone} } // decompressResponseBody 根据 Content-Encoding 解压响应体。 diff --git a/backend/internal/repository/http_upstream_redirect_test.go b/backend/internal/repository/http_upstream_redirect_test.go new file mode 100644 index 000000000..d1887d6ff --- /dev/null +++ b/backend/internal/repository/http_upstream_redirect_test.go @@ -0,0 +1,40 @@ +package repository + +import ( + "context" + "errors" + "net/http" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/service" +) + +func TestHTTPClientForUpstreamRequestDisablesRedirectsWithoutMutatingSharedClient(t *testing.T) { + shared := &http.Client{} + plainReq, err := http.NewRequestWithContext(context.Background(), http.MethodGet, "https://example.com", nil) + if err != nil { + t.Fatal(err) + } + if got := httpClientForUpstreamRequest(shared, plainReq); got != shared { + t.Fatal("ordinary request should reuse the shared client") + } + + secureCtx := service.WithHTTPUpstreamRedirectsDisabled(context.Background()) + secureReq, err := http.NewRequestWithContext(secureCtx, http.MethodGet, "https://example.com", nil) + if err != nil { + t.Fatal(err) + } + got := httpClientForUpstreamRequest(shared, secureReq) + if got == shared { + t.Fatal("redirect-disabled request should use a shallow client copy") + } + if shared.CheckRedirect != nil { + t.Fatal("shared client redirect policy must remain unchanged") + } + if got.CheckRedirect == nil { + t.Fatal("redirect-disabled client should install a redirect policy") + } + if err := got.CheckRedirect(secureReq, nil); !errors.Is(err, http.ErrUseLastResponse) { + t.Fatalf("CheckRedirect error = %v, want http.ErrUseLastResponse", err) + } +} diff --git a/backend/internal/repository/http_upstream_test.go b/backend/internal/repository/http_upstream_test.go index 5990e27c8..cda902872 100644 --- a/backend/internal/repository/http_upstream_test.go +++ b/backend/internal/repository/http_upstream_test.go @@ -1,6 +1,7 @@ package repository import ( + "errors" "io" "net/http" "sync/atomic" @@ -48,6 +49,9 @@ func (s *HTTPUpstreamSuite) TestDefaultResponseHeaderTimeout() { transport, ok := entry.client.Transport.(*http.Transport) require.True(s.T(), ok, "expected *http.Transport") require.Zero(s.T(), transport.ResponseHeaderTimeout, "ResponseHeaderTimeout mismatch") + require.NotNil(s.T(), transport.DialContext, "cold TCP dial must be bounded") + require.Equal(s.T(), defaultUpstreamTLSHandshakeTimeout, transport.TLSHandshakeTimeout, "TLS handshake timeout mismatch") + require.Equal(s.T(), defaultUpstreamExpectContinueTimeout, transport.ExpectContinueTimeout, "Expect-Continue timeout mismatch") } func (s *HTTPUpstreamSuite) TestNilConfigResponseHeaderTimeoutFallback() { @@ -177,6 +181,37 @@ func (s *HTTPUpstreamSuite) TestDo_EmptyProxy_UsesDirect() { require.Equal(s.T(), "direct-empty", string(b)) } +func (s *HTTPUpstreamSuite) TestDo_EOFWithoutCloseReleasesCapacityForEviction() { + upstream := newLocalTestServer(s.T(), http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = io.WriteString(w, "ok") + })) + s.T().Cleanup(upstream.Close) + + s.cfg.Gateway = config.GatewayConfig{ + ConnectionPoolIsolation: config.ConnectionPoolIsolationAccount, + MaxUpstreamClients: 1, + } + svc := s.newService() + + firstReq, err := http.NewRequest(http.MethodGet, upstream.URL, nil) + require.NoError(s.T(), err) + firstResp, err := svc.Do(firstReq, "", 1, 1) + require.NoError(s.T(), err) + _, err = io.ReadAll(firstResp.Body) + require.NoError(s.T(), err) + require.Len(s.T(), svc.clients, 1) + for _, entry := range svc.clients { + require.Zero(s.T(), atomic.LoadInt64(&entry.inFlight), "读取到 EOF 后应释放客户端占用") + } + + secondReq, err := http.NewRequest(http.MethodGet, upstream.URL, nil) + require.NoError(s.T(), err) + secondResp, err := svc.Do(secondReq, "", 2, 1) + require.NoError(s.T(), err, "已完成的首个请求不应阻止缓存淘汰") + require.NoError(s.T(), secondResp.Body.Close()) + require.Len(s.T(), svc.clients, 1, "缓存应保持在配置上限内") +} + // TestAccountIsolation_DifferentAccounts 测试账户隔离模式 // 验证不同账户使用独立的连接池 func (s *HTTPUpstreamSuite) TestAccountIsolation_DifferentAccounts() { @@ -286,6 +321,63 @@ func (s *HTTPUpstreamSuite) TestIdleTTLDoesNotEvictActive() { require.True(s.T(), hasEntry(svc, entry1), "有活跃请求时不应回收") } +func (s *HTTPUpstreamSuite) TestTrackedBodyReleasesOnEOFWithoutClose() { + var released atomic.Int64 + body := wrapTrackedBody(io.NopCloser(&singleReadEOFReader{data: []byte("ok")}), func() { + released.Add(1) + }) + + data, err := io.ReadAll(body) + require.NoError(s.T(), err) + require.Equal(s.T(), "ok", string(data)) + require.EqualValues(s.T(), 1, released.Load(), "EOF 应立即释放 inFlight") + + require.NoError(s.T(), body.Close()) + require.EqualValues(s.T(), 1, released.Load(), "Close 不应重复释放 inFlight") +} + +func (s *HTTPUpstreamSuite) TestTrackedBodyReleasesOnTerminalReadError() { + var released atomic.Int64 + body := wrapTrackedBody(io.NopCloser(errorReader{}), func() { + released.Add(1) + }) + + _, err := body.Read(make([]byte, 1)) + require.ErrorIs(s.T(), err, errTrackedBodyTest) + require.EqualValues(s.T(), 1, released.Load(), "终态读取错误应释放 inFlight") +} + +func (s *HTTPUpstreamSuite) TestWrapTrackedNilBodyReleasesImmediately() { + var released atomic.Int64 + body := wrapTrackedBody(nil, func() { + released.Add(1) + }) + + require.Nil(s.T(), body) + require.EqualValues(s.T(), 1, released.Load(), "空响应体不应泄漏 inFlight") +} + +type singleReadEOFReader struct { + data []byte +} + +func (r *singleReadEOFReader) Read(p []byte) (int, error) { + if len(r.data) == 0 { + return 0, io.EOF + } + n := copy(p, r.data) + r.data = r.data[n:] + return n, io.EOF +} + +var errTrackedBodyTest = errors.New("tracked body test error") + +type errorReader struct{} + +func (errorReader) Read([]byte) (int, error) { + return 0, errTrackedBodyTest +} + // TestHTTPUpstreamSuite 运行测试套件 func TestHTTPUpstreamSuite(t *testing.T) { suite.Run(t, new(HTTPUpstreamSuite)) diff --git a/backend/internal/repository/migrations_runner.go b/backend/internal/repository/migrations_runner.go index 36aab693c..b087da06f 100644 --- a/backend/internal/repository/migrations_runner.go +++ b/backend/internal/repository/migrations_runner.go @@ -51,6 +51,7 @@ CREATE TABLE IF NOT EXISTS atlas_schema_revisions ( // 任何稳定的 int64 值都可以，只要不与同一数据库中的其他锁冲突即可。 const migrationsAdvisoryLockID int64 = 694208311321144027 const migrationsLockRetryInterval = 500 * time.Millisecond +const migrationsUnlockTimeout = 5 * time.Second const nonTransactionalMigrationSuffix = "_notx.sql" const paymentOrdersOutTradeNoUniqueMigration = "120_enforce_payment_orders_out_trade_no_unique_notx.sql" const paymentOrdersOutTradeNoUniqueIndex = "paymentorder_out_trade_no_unique" @@ -59,6 +60,8 @@ const openAIOwnedAccountOrgIdentityUniqueMigration = "168_openai_owned_account_o const accountShareSeatCostQueryIndexesMigration = "208_account_share_seat_cost_query_indexes_notx.sql" const latestAPIKeyIPIndexMigration = "212_add_usage_logs_api_key_latest_ip_index_notx.sql" const latestAPIKeyIPIndex = "idx_usage_logs_api_key_latest_ip" +const usageLogImageInputTokensMigration = "216_usage_log_image_input_tokens.sql" +const openAIOwnedAgentIdentityUniqueMigration = "217_openai_owned_agent_identity_unique_notx.sql" const accountShareSeatCostAutoIndexMaxRows int64 = 5_000_000 const accountShareSeatCostAutoIndexMaxTableBytes int64 = 8 << 30 @@ -67,17 +70,30 @@ type migrationCatalogName struct { name string } +type migrationQueryExecutor interface { + ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) + QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) + QueryRowContext(ctx context.Context, query string, args ...any) *sql.Row +} + +type migrationDatabase interface { + migrationQueryExecutor + BeginTx(ctx context.Context, opts *sql.TxOptions) (*sql.Tx, error) +} + type migrationIndexKeyRequirement struct { column string expressionCanonical string resultType migrationCatalogName operatorClass migrationCatalogName + collation migrationCatalogName } type migrationIndexRequirement struct { index migrationCatalogName table migrationCatalogName accessMethod string + unique bool keys []migrationIndexKeyRequirement includeColumns []string predicateCanonical string @@ -124,6 +140,22 @@ var migrationTimestamptzKey = migrationIndexKeyRequirement{ operatorClass: migrationCatalogName{schema: "pg_catalog", name: "timestamptz_ops"}, } +var migrationTextKey = migrationIndexKeyRequirement{ + resultType: migrationCatalogName{schema: "pg_catalog", name: "text"}, + operatorClass: migrationCatalogName{schema: "pg_catalog", name: "text_ops"}, + collation: migrationCatalogName{schema: "pg_catalog", name: "default"}, +} + +func migrationColumnKeyRequirement(column string, template migrationIndexKeyRequirement) migrationIndexKeyRequirement { + template.column = column + return template +} + +func migrationExpressionKeyRequirement(expression string, template migrationIndexKeyRequirement) migrationIndexKeyRequirement { + template.expressionCanonical = canonicalizeMigrationIndexExpression(expression) + return template +} + const accountShareSeatReasonPredicateCanonical = "reasonIN['account_share_mode_seat_prepay','account_share_mode_seat_refund','account_share_mode_seat_waiver_refund']" const accountShareSeatMembershipExpressionCanonical = "NULLIFmetadata->>'membership_id',''" @@ -194,11 +226,20 @@ var ownedAccountIdentityUniqueIndexes = []string{ } var ownedAccountIdentityUniqueIndexSet = map[string]struct{}{ - "idx_accounts_owned_openai_chatgpt_account_id_uniq": {}, - "idx_accounts_owned_openai_chatgpt_user_id_uniq": {}, - "idx_accounts_owned_anthropic_org_account_uniq": {}, - "idx_accounts_owned_gemini_project_uniq": {}, - "idx_accounts_owned_antigravity_project_uniq": {}, + "idx_accounts_owned_openai_chatgpt_account_id_uniq": {}, + "idx_accounts_owned_openai_chatgpt_user_id_uniq": {}, + "idx_accounts_owned_openai_org_user_uniq": {}, + "idx_accounts_owned_openai_org_account_uniq": {}, + "idx_accounts_owned_openai_legacy_user_uniq": {}, + "idx_accounts_owned_openai_legacy_account_uniq": {}, + "idx_accounts_owned_openai_org_user_v2_uniq": {}, + "idx_accounts_owned_openai_org_account_v2_uniq": {}, + "idx_accounts_owned_openai_legacy_user_v2_uniq": {}, + "idx_accounts_owned_openai_legacy_account_v2_uniq": {}, + "idx_accounts_owned_openai_agent_identity_team_uniq": {}, + "idx_accounts_owned_anthropic_org_account_uniq": {}, + "idx_accounts_owned_gemini_project_uniq": {}, + "idx_accounts_owned_antigravity_project_uniq": {}, } var openAIOwnedAccountOrgIdentityUniqueIndexes = []string{ @@ -208,6 +249,93 @@ var openAIOwnedAccountOrgIdentityUniqueIndexes = []string{ "idx_accounts_owned_openai_legacy_account_uniq", } +const ( + openAIOwnedOrganizationExpression = `lower(NULLIF(btrim(credentials->>'organization_id'), ''))` + openAIOwnedUserExpression = `NULLIF(btrim(credentials->>'chatgpt_user_id'), '')` + openAIOwnedAccountExpression = `NULLIF(btrim(credentials->>'chatgpt_account_id'), '')` + openAIOwnedOAuthV2BasePredicate = `deleted_at IS NULL + AND owner_user_id IS NOT NULL + AND platform = 'openai' + AND type = 'oauth' + AND COALESCE(lower(NULLIF(btrim(credentials->>'auth_mode'), '')), '') <> 'agentidentity'` + openAIOwnedAgentIdentityPredicate = `deleted_at IS NULL + AND owner_user_id IS NOT NULL + AND platform = 'openai' + AND type = 'oauth' + AND lower(NULLIF(btrim(credentials->>'auth_mode'), '')) = 'agentidentity' + AND NULLIF(btrim(credentials->>'chatgpt_account_id'), '') IS NOT NULL` +) + +var openAIOwnedAgentIdentityUniqueIndexRequirements = []migrationIndexRequirement{ + { + index: migrationCatalogName{schema: "public", name: "idx_accounts_owned_openai_org_user_v2_uniq"}, + table: migrationCatalogName{schema: "public", name: "accounts"}, + accessMethod: "btree", + unique: true, + keys: []migrationIndexKeyRequirement{ + migrationColumnKeyRequirement("owner_user_id", migrationInt8Key), + migrationExpressionKeyRequirement(openAIOwnedOrganizationExpression, migrationTextKey), + migrationExpressionKeyRequirement(openAIOwnedUserExpression, migrationTextKey), + }, + predicateCanonical: canonicalizeMigrationIndexExpression(openAIOwnedOAuthV2BasePredicate + ` + AND NULLIF(btrim(credentials->>'organization_id'), '') IS NOT NULL + AND NULLIF(btrim(credentials->>'chatgpt_user_id'), '') IS NOT NULL`), + }, + { + index: migrationCatalogName{schema: "public", name: "idx_accounts_owned_openai_org_account_v2_uniq"}, + table: migrationCatalogName{schema: "public", name: "accounts"}, + accessMethod: "btree", + unique: true, + keys: []migrationIndexKeyRequirement{ + migrationColumnKeyRequirement("owner_user_id", migrationInt8Key), + migrationExpressionKeyRequirement(openAIOwnedOrganizationExpression, migrationTextKey), + migrationExpressionKeyRequirement(openAIOwnedAccountExpression, migrationTextKey), + }, + predicateCanonical: canonicalizeMigrationIndexExpression(openAIOwnedOAuthV2BasePredicate + ` + AND NULLIF(btrim(credentials->>'organization_id'), '') IS NOT NULL + AND NULLIF(btrim(credentials->>'chatgpt_user_id'), '') IS NULL + AND NULLIF(btrim(credentials->>'chatgpt_account_id'), '') IS NOT NULL`), + }, + { + index: migrationCatalogName{schema: "public", name: "idx_accounts_owned_openai_legacy_user_v2_uniq"}, + table: migrationCatalogName{schema: "public", name: "accounts"}, + accessMethod: "btree", + unique: true, + keys: []migrationIndexKeyRequirement{ + migrationColumnKeyRequirement("owner_user_id", migrationInt8Key), + migrationExpressionKeyRequirement(openAIOwnedUserExpression, migrationTextKey), + }, + predicateCanonical: canonicalizeMigrationIndexExpression(openAIOwnedOAuthV2BasePredicate + ` + AND NULLIF(btrim(credentials->>'organization_id'), '') IS NULL + AND NULLIF(btrim(credentials->>'chatgpt_user_id'), '') IS NOT NULL`), + }, + { + index: migrationCatalogName{schema: "public", name: "idx_accounts_owned_openai_legacy_account_v2_uniq"}, + table: migrationCatalogName{schema: "public", name: "accounts"}, + accessMethod: "btree", + unique: true, + keys: []migrationIndexKeyRequirement{ + migrationColumnKeyRequirement("owner_user_id", migrationInt8Key), + migrationExpressionKeyRequirement(openAIOwnedAccountExpression, migrationTextKey), + }, + predicateCanonical: canonicalizeMigrationIndexExpression(openAIOwnedOAuthV2BasePredicate + ` + AND NULLIF(btrim(credentials->>'organization_id'), '') IS NULL + AND NULLIF(btrim(credentials->>'chatgpt_user_id'), '') IS NULL + AND NULLIF(btrim(credentials->>'chatgpt_account_id'), '') IS NOT NULL`), + }, + { + index: migrationCatalogName{schema: "public", name: "idx_accounts_owned_openai_agent_identity_team_uniq"}, + table: migrationCatalogName{schema: "public", name: "accounts"}, + accessMethod: "btree", + unique: true, + keys: []migrationIndexKeyRequirement{ + migrationColumnKeyRequirement("owner_user_id", migrationInt8Key), + migrationExpressionKeyRequirement(openAIOwnedAccountExpression, migrationTextKey), + }, + predicateCanonical: canonicalizeMigrationIndexExpression(openAIOwnedAgentIdentityPredicate), + }, +} + type migrationChecksumCompatibilityRule struct { fileChecksum string acceptedDBChecksum map[string]struct{} @@ -280,7 +408,17 @@ func applyMigrationsFS(ctx context.Context, db *sql.DB, fsys fs.FS) error { if db == nil { return errors.New("nil sql db") } + conn, err := db.Conn(ctx) + if err != nil { + return fmt.Errorf("pin migration database connection: %w", err) + } + defer func() { + _ = conn.Close() + }() + return applyMigrationsOnConnectionFS(ctx, conn, fsys) +} +func applyMigrationsOnConnectionFS(ctx context.Context, db migrationDatabase, fsys fs.FS) error { // 获取分布式锁，确保多实例部署时只有一个实例执行迁移。 // 这是 PostgreSQL 特有的 Advisory Lock 机制。 if err := pgAdvisoryLock(ctx, db); err != nil { @@ -288,8 +426,11 @@ func applyMigrationsFS(ctx context.Context, db *sql.DB, fsys fs.FS) error { } defer func() { // 无论迁移是否成功，都要释放锁。 - // 使用 context.Background() 确保即使原 ctx 已取消也能释放锁。 - _ = pgAdvisoryUnlock(context.Background(), db) + // 原迁移上下文可能已取消，因此使用独立的短超时收尾；即使 + // 显式解锁失败，随后关闭固定连接也会释放 session advisory lock。 + unlockCtx, cancel := context.WithTimeout(context.Background(), migrationsUnlockTimeout) + defer cancel() + _ = pgAdvisoryUnlock(unlockCtx, db) }() // 创建迁移记录表（如果不存在）。 @@ -369,14 +510,24 @@ func applyMigrationsFS(ctx context.Context, db *sql.DB, fsys fs.FS) error { // *_notx.sql：用于 CREATE/DROP INDEX CONCURRENTLY 场景，必须非事务执行。 // 逐条语句执行，避免将多条 CONCURRENTLY 语句放入同一个隐式事务块。 statements := splitSQLStatements(content) + verifiedAgentIdentityIndexesBeforeDrop := false for i, stmt := range statements { trimmed := strings.TrimSpace(stmt) if trimmed == "" { continue } - if stripSQLLineComment(trimmed) == "" { + statementWithoutComments := stripSQLLineComment(trimmed) + if statementWithoutComments == "" { continue } + if name == openAIOwnedAgentIdentityUniqueMigration && + !verifiedAgentIdentityIndexesBeforeDrop && + strings.HasPrefix(strings.ToUpper(statementWithoutComments), "DROP INDEX CONCURRENTLY") { + if err := verifyNonTransactionalMigrationResult(ctx, db, name); err != nil { + return fmt.Errorf("verify migration %s before dropping protected indexes: %w", name, err) + } + verifiedAgentIdentityIndexesBeforeDrop = true + } if _, err := db.ExecContext(ctx, trimmed); err != nil { return fmt.Errorf("apply migration %s (non-tx statement %d): %w", name, i+1, err) } @@ -395,12 +546,24 @@ func applyMigrationsFS(ctx context.Context, db *sql.DB, fsys fs.FS) error { if err != nil { return fmt.Errorf("begin migration %s: %w", name, err) } + if name == usageLogImageInputTokensMigration { + // Listing pg_temp explicitly after public prevents PostgreSQL from + // implicitly searching a temporary schema before public for relations. + if _, err := tx.ExecContext(ctx, "SET LOCAL search_path = pg_catalog, public, pg_temp"); err != nil { + _ = tx.Rollback() + return fmt.Errorf("pin migration %s search_path: %w", name, err) + } + } // 执行迁移 SQL if _, err := tx.ExecContext(ctx, content); err != nil { _ = tx.Rollback() return fmt.Errorf("apply migration %s: %w", name, err) } + if err := verifyTransactionalMigrationResult(ctx, tx, name); err != nil { + _ = tx.Rollback() + return fmt.Errorf("verify migration %s: %w", name, err) + } // 记录迁移已完成，保存文件名和校验和 if _, err := tx.ExecContext(ctx, "INSERT INTO schema_migrations (filename, checksum) VALUES ($1, $2)", name, checksum); err != nil { @@ -418,7 +581,73 @@ func applyMigrationsFS(ctx context.Context, db *sql.DB, fsys fs.FS) error { return nil } -func prepareNonTransactionalMigration(ctx context.Context, db *sql.DB, name string) error { +func verifyTransactionalMigrationResult(ctx context.Context, db migrationQueryExecutor, name string) error { + if name != usageLogImageInputTokensMigration { + return nil + } + + rows, err := db.QueryContext(ctx, ` + SELECT + a.attname, + pg_catalog.format_type(a.atttypid, a.atttypmod), + a.attnotnull, + COALESCE(pg_catalog.pg_get_expr(d.adbin, d.adrelid), '') + FROM pg_catalog.pg_attribute AS a + JOIN pg_catalog.pg_class AS c ON c.oid = a.attrelid + JOIN pg_catalog.pg_namespace AS n ON n.oid = c.relnamespace + LEFT JOIN pg_catalog.pg_attrdef AS d + ON d.adrelid = a.attrelid + AND d.adnum = a.attnum + WHERE n.nspname = 'public' + AND c.relname = 'usage_logs' + AND c.relkind IN ('r', 'p') + AND a.attname IN ('image_input_tokens', 'image_input_cost') + AND a.attnum > 0 + AND NOT a.attisdropped + ORDER BY a.attname + `) + if err != nil { + return fmt.Errorf("inspect public.usage_logs image input columns: %w", err) + } + defer func() { _ = rows.Close() }() + + expectedTypes := map[string]string{ + "image_input_cost": "numeric(20,10)", + "image_input_tokens": "integer", + } + seen := make(map[string]struct{}, len(expectedTypes)) + for rows.Next() { + var columnName string + var dataType string + var notNull bool + var defaultExpression string + if err := rows.Scan(&columnName, &dataType, ¬Null, &defaultExpression); err != nil { + return fmt.Errorf("scan public.usage_logs image input column: %w", err) + } + expectedType, ok := expectedTypes[columnName] + if !ok || dataType != expectedType || !notNull || strings.TrimSpace(defaultExpression) != "0" { + return fmt.Errorf( + "public.usage_logs.%s has unexpected definition (type=%s not_null=%t default=%q)", + columnName, + dataType, + notNull, + defaultExpression, + ) + } + seen[columnName] = struct{}{} + } + if err := rows.Err(); err != nil { + return fmt.Errorf("iterate public.usage_logs image input columns: %w", err) + } + for columnName := range expectedTypes { + if _, ok := seen[columnName]; !ok { + return fmt.Errorf("public.usage_logs.%s is missing", columnName) + } + } + return nil +} + +func prepareNonTransactionalMigration(ctx context.Context, db migrationDatabase, name string) error { switch name { case paymentOrdersOutTradeNoUniqueMigration: return preparePaymentOrdersOutTradeNoUniqueMigration(ctx, db) @@ -426,6 +655,8 @@ func prepareNonTransactionalMigration(ctx context.Context, db *sql.DB, name stri return prepareOwnedAccountIdentityUniqueMigration(ctx, db) case openAIOwnedAccountOrgIdentityUniqueMigration: return prepareOpenAIOwnedAccountOrgIdentityUniqueMigration(ctx, db) + case openAIOwnedAgentIdentityUniqueMigration: + return prepareOpenAIOwnedAgentIdentityUniqueMigration(ctx, db) case accountShareSeatCostQueryIndexesMigration: return prepareAccountShareSeatCostQueryIndexesMigration(ctx, db) case latestAPIKeyIPIndexMigration: @@ -435,7 +666,7 @@ func prepareNonTransactionalMigration(ctx context.Context, db *sql.DB, name stri } } -func prepareLatestAPIKeyIPIndexMigration(ctx context.Context, db *sql.DB) error { +func prepareLatestAPIKeyIPIndexMigration(ctx context.Context, db migrationDatabase) error { invalid, err := indexIsInvalid(ctx, db, latestAPIKeyIPIndex) if err != nil { return fmt.Errorf("check invalid index %s: %w", latestAPIKeyIPIndex, err) @@ -449,7 +680,17 @@ func prepareLatestAPIKeyIPIndexMigration(ctx context.Context, db *sql.DB) error return nil } -func verifyNonTransactionalMigrationResult(ctx context.Context, db *sql.DB, name string) error { +func verifyNonTransactionalMigrationResult(ctx context.Context, db migrationDatabase, name string) error { + if name == openAIOwnedAgentIdentityUniqueMigration { + matches, err := indexesMatchRequirements(ctx, db, openAIOwnedAgentIdentityUniqueIndexRequirements) + if err != nil { + return fmt.Errorf("verify OpenAI owned Agent Identity unique index definitions: %w", err) + } + if !matches { + return errors.New("one or more OpenAI owned Agent Identity unique indexes are missing, invalid, or do not match migration 217") + } + return nil + } if name != accountShareSeatCostQueryIndexesMigration { return nil } @@ -463,7 +704,7 @@ func verifyNonTransactionalMigrationResult(ctx context.Context, db *sql.DB, name return nil } -func prepareAccountShareSeatCostQueryIndexesMigration(ctx context.Context, db *sql.DB) error { +func prepareAccountShareSeatCostQueryIndexesMigration(ctx context.Context, db migrationDatabase) error { ready, err := indexesMatchRequirements(ctx, db, accountShareSeatCostIndexRequirements[:2]) if err != nil { return fmt.Errorf("check account-share seat-cost ledger indexes: %w", err) @@ -485,7 +726,7 @@ func prepareAccountShareSeatCostQueryIndexesMigration(ctx context.Context, db *s return prepareIndexesForRetry(ctx, db, accountShareSeatCostIndexRequirements) } -func prepareIndexesForRetry(ctx context.Context, db *sql.DB, requirements []migrationIndexRequirement) error { +func prepareIndexesForRetry(ctx context.Context, db migrationDatabase, requirements []migrationIndexRequirement) error { for _, requirement := range requirements { indexName := requirement.index.name invalid, err := indexIsInvalid(ctx, db, indexName) @@ -510,7 +751,7 @@ func quoteMigrationCatalogName(name migrationCatalogName) string { return quoteIdentifier(name.schema) + "." + quoteIdentifier(name.name) } -func preparePaymentOrdersOutTradeNoUniqueMigration(ctx context.Context, db *sql.DB) error { +func preparePaymentOrdersOutTradeNoUniqueMigration(ctx context.Context, db migrationDatabase) error { duplicates, err := findDuplicatePaymentOrderOutTradeNos(ctx, db) if err != nil { return fmt.Errorf("precheck duplicate out_trade_no: %w", err) @@ -537,7 +778,7 @@ func preparePaymentOrdersOutTradeNoUniqueMigration(ctx context.Context, db *sql. return nil } -func prepareOwnedAccountIdentityUniqueMigration(ctx context.Context, db *sql.DB) error { +func prepareOwnedAccountIdentityUniqueMigration(ctx context.Context, db migrationDatabase) error { duplicates, err := findDuplicateOwnedAccountIdentities(ctx, db) if err != nil { return fmt.Errorf("precheck duplicate owned account identities: %w", err) @@ -565,7 +806,7 @@ func prepareOwnedAccountIdentityUniqueMigration(ctx context.Context, db *sql.DB) return nil } -func prepareOpenAIOwnedAccountOrgIdentityUniqueMigration(ctx context.Context, db *sql.DB) error { +func prepareOpenAIOwnedAccountOrgIdentityUniqueMigration(ctx context.Context, db migrationDatabase) error { duplicates, err := findDuplicateOpenAIOwnedAccountOrgIdentities(ctx, db) if err != nil { return fmt.Errorf("precheck duplicate OpenAI owned account org identities: %w", err) @@ -593,7 +834,55 @@ func prepareOpenAIOwnedAccountOrgIdentityUniqueMigration(ctx context.Context, db return nil } -func findDuplicateOwnedAccountIdentities(ctx context.Context, db *sql.DB) ([]string, error) { +func prepareOpenAIOwnedAgentIdentityUniqueMigration(ctx context.Context, db migrationDatabase) error { + duplicates, err := findDuplicateOpenAIOwnedAgentIdentityMigrationIdentities(ctx, db) + if err != nil { + return fmt.Errorf("precheck duplicate OpenAI owned Agent Identity migration identities: %w", err) + } + if len(duplicates) > 0 { + return fmt.Errorf( + "duplicate OpenAI owned Agent Identity migration identities block %s; remediate duplicates before retrying: %s", + openAIOwnedAgentIdentityUniqueMigration, + strings.Join(duplicates, ", "), + ) + } + + for _, requirement := range openAIOwnedAgentIdentityUniqueIndexRequirements { + invalid, err := indexIsInvalid(ctx, db, requirement.index.name) + if err != nil { + return fmt.Errorf("check invalid index %s: %w", requirement.index.name, err) + } + if invalid { + qualifiedIndexName := quoteMigrationCatalogName(requirement.index) + if _, err := db.ExecContext(ctx, fmt.Sprintf("DROP INDEX CONCURRENTLY IF EXISTS %s", qualifiedIndexName)); err != nil { + return fmt.Errorf("drop invalid index %s: %w", requirement.index.name, err) + } + continue + } + + state, err := loadMigrationIndexCatalogState(ctx, db, requirement.index) + if err != nil { + return fmt.Errorf("inspect existing index %s: %w", requirement.index.name, err) + } + if !state.exists || migrationIndexMatchesRequirement(state, requirement) { + continue + } + if state.ready && state.valid && state.live { + return fmt.Errorf( + "existing index %s has a valid but unexpected definition; refuse to continue %s until the conflicting index is reviewed", + requirement.index.name, + openAIOwnedAgentIdentityUniqueMigration, + ) + } + qualifiedIndexName := quoteMigrationCatalogName(requirement.index) + if _, err := db.ExecContext(ctx, fmt.Sprintf("DROP INDEX CONCURRENTLY IF EXISTS %s", qualifiedIndexName)); err != nil { + return fmt.Errorf("drop invalid index %s: %w", requirement.index.name, err) + } + } + return nil +} + +func findDuplicateOwnedAccountIdentities(ctx context.Context, db migrationDatabase) ([]string, error) { rows, err := db.QueryContext(ctx, ` WITH identities AS ( SELECT @@ -601,7 +890,7 @@ func findDuplicateOwnedAccountIdentities(ctx context.Context, db *sql.DB) ([]str 'openai.chatgpt_account_id' AS identity_name, NULLIF(BTRIM(credentials->>'chatgpt_account_id'), '') AS identity_value, id - FROM accounts + FROM public.accounts WHERE deleted_at IS NULL AND owner_user_id IS NOT NULL AND platform = 'openai' @@ -615,7 +904,7 @@ func findDuplicateOwnedAccountIdentities(ctx context.Context, db *sql.DB) ([]str 'openai.chatgpt_user_id' AS identity_name, NULLIF(BTRIM(credentials->>'chatgpt_user_id'), '') AS identity_value, id - FROM accounts + FROM public.accounts WHERE deleted_at IS NULL AND owner_user_id IS NOT NULL AND platform = 'openai' @@ -631,7 +920,7 @@ func findDuplicateOwnedAccountIdentities(ctx context.Context, db *sql.DB) ([]str '|' || LOWER(COALESCE(NULLIF(BTRIM(extra->>'account_uuid'), ''), NULLIF(BTRIM(credentials->>'account_uuid'), ''))) AS identity_value, id - FROM accounts + FROM public.accounts WHERE deleted_at IS NULL AND owner_user_id IS NOT NULL AND platform = 'anthropic' @@ -648,7 +937,7 @@ func findDuplicateOwnedAccountIdentities(ctx context.Context, db *sql.DB) ([]str '|' || LOWER(NULLIF(BTRIM(credentials->>'project_id'), '')) AS identity_value, id - FROM accounts + FROM public.accounts WHERE deleted_at IS NULL AND owner_user_id IS NOT NULL AND platform = 'gemini' @@ -662,7 +951,7 @@ func findDuplicateOwnedAccountIdentities(ctx context.Context, db *sql.DB) ([]str 'antigravity.project_id' AS identity_name, LOWER(NULLIF(BTRIM(credentials->>'project_id'), '')) AS identity_value, id - FROM accounts + FROM public.accounts WHERE deleted_at IS NULL AND owner_user_id IS NOT NULL AND platform = 'antigravity' @@ -710,7 +999,7 @@ func findDuplicateOwnedAccountIdentities(ctx context.Context, db *sql.DB) ([]str return duplicates, nil } -func findDuplicateOpenAIOwnedAccountOrgIdentities(ctx context.Context, db *sql.DB) ([]string, error) { +func findDuplicateOpenAIOwnedAccountOrgIdentities(ctx context.Context, db migrationDatabase) ([]string, error) { rows, err := db.QueryContext(ctx, ` WITH identities AS ( SELECT @@ -818,7 +1107,135 @@ func findDuplicateOpenAIOwnedAccountOrgIdentities(ctx context.Context, db *sql.D return duplicates, nil } -func findDuplicatePaymentOrderOutTradeNos(ctx context.Context, db *sql.DB) ([]string, error) { +func findDuplicateOpenAIOwnedAgentIdentityMigrationIdentities(ctx context.Context, db migrationDatabase) ([]string, error) { + rows, err := db.QueryContext(ctx, ` + WITH identities AS ( + SELECT + owner_user_id, + 'openai.org_user_v2' AS identity_name, + LOWER(NULLIF(BTRIM(credentials->>'organization_id'), '')) AS identity_value_1, + NULLIF(BTRIM(credentials->>'chatgpt_user_id'), '') AS identity_value_2, + id + FROM public.accounts + WHERE deleted_at IS NULL + AND owner_user_id IS NOT NULL + AND platform = 'openai' + AND type = 'oauth' + AND COALESCE(LOWER(NULLIF(BTRIM(credentials->>'auth_mode'), '')), '') <> 'agentidentity' + AND NULLIF(BTRIM(credentials->>'organization_id'), '') IS NOT NULL + AND NULLIF(BTRIM(credentials->>'chatgpt_user_id'), '') IS NOT NULL + + UNION ALL + + SELECT + owner_user_id, + 'openai.org_account_v2' AS identity_name, + LOWER(NULLIF(BTRIM(credentials->>'organization_id'), '')) AS identity_value_1, + NULLIF(BTRIM(credentials->>'chatgpt_account_id'), '') AS identity_value_2, + id + FROM public.accounts + WHERE deleted_at IS NULL + AND owner_user_id IS NOT NULL + AND platform = 'openai' + AND type = 'oauth' + AND COALESCE(LOWER(NULLIF(BTRIM(credentials->>'auth_mode'), '')), '') <> 'agentidentity' + AND NULLIF(BTRIM(credentials->>'organization_id'), '') IS NOT NULL + AND NULLIF(BTRIM(credentials->>'chatgpt_user_id'), '') IS NULL + AND NULLIF(BTRIM(credentials->>'chatgpt_account_id'), '') IS NOT NULL + + UNION ALL + + SELECT + owner_user_id, + 'openai.legacy_user_v2' AS identity_name, + NULLIF(BTRIM(credentials->>'chatgpt_user_id'), '') AS identity_value_1, + NULL::text AS identity_value_2, + id + FROM public.accounts + WHERE deleted_at IS NULL + AND owner_user_id IS NOT NULL + AND platform = 'openai' + AND type = 'oauth' + AND COALESCE(LOWER(NULLIF(BTRIM(credentials->>'auth_mode'), '')), '') <> 'agentidentity' + AND NULLIF(BTRIM(credentials->>'organization_id'), '') IS NULL + AND NULLIF(BTRIM(credentials->>'chatgpt_user_id'), '') IS NOT NULL + + UNION ALL + + SELECT + owner_user_id, + 'openai.legacy_account_v2' AS identity_name, + NULLIF(BTRIM(credentials->>'chatgpt_account_id'), '') AS identity_value_1, + NULL::text AS identity_value_2, + id + FROM public.accounts + WHERE deleted_at IS NULL + AND owner_user_id IS NOT NULL + AND platform = 'openai' + AND type = 'oauth' + AND COALESCE(LOWER(NULLIF(BTRIM(credentials->>'auth_mode'), '')), '') <> 'agentidentity' + AND NULLIF(BTRIM(credentials->>'organization_id'), '') IS NULL + AND NULLIF(BTRIM(credentials->>'chatgpt_user_id'), '') IS NULL + AND NULLIF(BTRIM(credentials->>'chatgpt_account_id'), '') IS NOT NULL + + UNION ALL + + SELECT + owner_user_id, + 'openai.agent_identity_team' AS identity_name, + NULLIF(BTRIM(credentials->>'chatgpt_account_id'), '') AS identity_value_1, + NULL::text AS identity_value_2, + id + FROM public.accounts + WHERE deleted_at IS NULL + AND owner_user_id IS NOT NULL + AND platform = 'openai' + AND type = 'oauth' + AND LOWER(NULLIF(BTRIM(credentials->>'auth_mode'), '')) = 'agentidentity' + AND NULLIF(BTRIM(credentials->>'chatgpt_account_id'), '') IS NOT NULL + ) + SELECT + identity_name, + owner_user_id, + COUNT(*) AS duplicate_count, + ARRAY_TO_STRING((ARRAY_AGG(id ORDER BY id))[1:5], ',') AS sample_ids + FROM identities + GROUP BY identity_name, owner_user_id, identity_value_1, identity_value_2 + HAVING COUNT(*) > 1 + ORDER BY duplicate_count DESC, identity_name, owner_user_id + LIMIT 10 + `) + if err != nil { + return nil, err + } + defer func() { + _ = rows.Close() + }() + + duplicates := make([]string, 0, 10) + for rows.Next() { + var identityName string + var ownerUserID int64 + var duplicateCount int + var sampleIDs string + if err := rows.Scan(&identityName, &ownerUserID, &duplicateCount, &sampleIDs); err != nil { + return nil, err + } + duplicates = append(duplicates, fmt.Sprintf( + "%s owner_user_id=%d count=%d sample_account_ids=%s", + identityName, + ownerUserID, + duplicateCount, + sampleIDs, + )) + } + if err := rows.Err(); err != nil { + return nil, err + } + return duplicates, nil +} + +func findDuplicatePaymentOrderOutTradeNos(ctx context.Context, db migrationDatabase) ([]string, error) { rows, err := db.QueryContext(ctx, ` SELECT out_trade_no, COUNT(*) AS duplicate_count FROM payment_orders @@ -850,7 +1267,7 @@ func findDuplicatePaymentOrderOutTradeNos(ctx context.Context, db *sql.DB) ([]st return duplicates, nil } -func indexIsInvalid(ctx context.Context, db *sql.DB, indexName string) (bool, error) { +func indexIsInvalid(ctx context.Context, db migrationDatabase, indexName string) (bool, error) { var invalid bool err := db.QueryRowContext(ctx, ` SELECT EXISTS ( @@ -866,7 +1283,7 @@ func indexIsInvalid(ctx context.Context, db *sql.DB, indexName string) (bool, er return invalid, err } -func indexesMatchRequirements(ctx context.Context, db *sql.DB, requirements []migrationIndexRequirement) (bool, error) { +func indexesMatchRequirements(ctx context.Context, db migrationDatabase, requirements []migrationIndexRequirement) (bool, error) { for _, requirement := range requirements { state, err := loadMigrationIndexCatalogState(ctx, db, requirement.index) if err != nil { @@ -879,7 +1296,7 @@ func indexesMatchRequirements(ctx context.Context, db *sql.DB, requirements []mi return true, nil } -func loadMigrationIndexCatalogState(ctx context.Context, db *sql.DB, index migrationCatalogName) (migrationIndexCatalogState, error) { +func loadMigrationIndexCatalogState(ctx context.Context, db migrationDatabase, index migrationCatalogName) (migrationIndexCatalogState, error) { var state migrationIndexCatalogState var keysJSON string var includeColumnsJSON string @@ -1001,7 +1418,7 @@ func migrationIndexMatchesRequirement(state migrationIndexCatalogState, requirem if state.table != requirement.table || state.accessMethod != requirement.accessMethod || state.relationKind != "i" || - state.unique || state.primary || state.exclusion || + state.unique != requirement.unique || state.primary || state.exclusion || !state.ready || !state.valid || !state.live || state.keyCount != len(requirement.keys) || state.attributeCount != len(requirement.keys)+len(requirement.includeColumns) || @@ -1016,7 +1433,8 @@ func migrationIndexMatchesRequirement(state migrationIndexCatalogState, requirem actual.TypeName != expected.resultType.name || actual.OpClassSchema != expected.operatorClass.schema || actual.OpClassName != expected.operatorClass.name || - actual.CollationSchema != "" || actual.CollationName != "" || + actual.CollationSchema != expected.collation.schema || + actual.CollationName != expected.collation.name || actual.OptionBits != 0 { return false } @@ -1087,7 +1505,7 @@ func isMigrationIndexCastContinuation(next byte) bool { (next >= 'a' && next <= 'z') } -func tableRowAndSizeEstimates(ctx context.Context, db *sql.DB, tableName string) (int64, int64, error) { +func tableRowAndSizeEstimates(ctx context.Context, db migrationDatabase, tableName string) (int64, int64, error) { var estimatedRows int64 var tableBytes int64 err := db.QueryRowContext(ctx, ` @@ -1101,7 +1519,7 @@ func tableRowAndSizeEstimates(ctx context.Context, db *sql.DB, tableName string) return estimatedRows, tableBytes, err } -func ensureAtlasBaselineAligned(ctx context.Context, db *sql.DB, fsys fs.FS) error { +func ensureAtlasBaselineAligned(ctx context.Context, db migrationDatabase, fsys fs.FS) error { hasLegacy, err := tableExists(ctx, db, "schema_migrations") if err != nil { return fmt.Errorf("check schema_migrations: %w", err) @@ -1142,7 +1560,7 @@ func ensureAtlasBaselineAligned(ctx context.Context, db *sql.DB, fsys fs.FS) err return nil } -func tableExists(ctx context.Context, db *sql.DB, tableName string) (bool, error) { +func tableExists(ctx context.Context, db migrationDatabase, tableName string) (bool, error) { var exists bool err := db.QueryRowContext(ctx, ` SELECT EXISTS ( @@ -1273,7 +1691,7 @@ func stripSQLLineComment(s string) string { // pgAdvisoryLock 获取 PostgreSQL Advisory Lock。 // Advisory Lock 是一种轻量级的锁机制，不与任何特定的数据库对象关联。 // 它非常适合用于应用层面的分布式锁场景，如迁移序列化。 -func pgAdvisoryLock(ctx context.Context, db *sql.DB) error { +func pgAdvisoryLock(ctx context.Context, db migrationDatabase) error { ticker := time.NewTicker(migrationsLockRetryInterval) defer ticker.Stop() @@ -1295,7 +1713,7 @@ func pgAdvisoryLock(ctx context.Context, db *sql.DB) error { // pgAdvisoryUnlock 释放 PostgreSQL Advisory Lock。 // 必须在获取锁后确保释放，否则会阻塞其他实例的迁移操作。 -func pgAdvisoryUnlock(ctx context.Context, db *sql.DB) error { +func pgAdvisoryUnlock(ctx context.Context, db migrationDatabase) error { _, err := db.ExecContext(ctx, "SELECT pg_advisory_unlock($1)", migrationsAdvisoryLockID) if err != nil { return fmt.Errorf("release migrations lock: %w", err) diff --git a/backend/internal/repository/migrations_runner_notx_test.go b/backend/internal/repository/migrations_runner_notx_test.go index e511a69bf..553e31ab8 100644 --- a/backend/internal/repository/migrations_runner_notx_test.go +++ b/backend/internal/repository/migrations_runner_notx_test.go @@ -8,6 +8,8 @@ import ( "testing" "testing/fstest" + "github.com/Wei-Shaw/sub2api/migrations" + sqlmock "github.com/DATA-DOG/go-sqlmock" "github.com/stretchr/testify/require" ) @@ -30,6 +32,8 @@ var migrationIndexCatalogColumns = []string{ "predicate", } +const openAIAgentIdentityDuplicatePrecheckPattern = `(?s)WITH identities AS.*identity_value_1.*identity_value_2.*FROM public\.accounts.*GROUP BY identity_name, owner_user_id, identity_value_1, identity_value_2` + func matchingMigrationIndexCatalogRows( t *testing.T, requirement migrationIndexRequirement, @@ -41,6 +45,7 @@ func matchingMigrationIndexCatalogRows( table: requirement.table, accessMethod: requirement.accessMethod, relationKind: "i", + unique: requirement.unique, ready: true, valid: true, live: true, @@ -53,26 +58,25 @@ func matchingMigrationIndexCatalogRows( definition := expected.column isExpression := expected.expressionCanonical != "" if isExpression { - definition = "(NULLIF((metadata ->> 'membership_id'::text), ''::text))::bigint" + definition = expected.expressionCanonical } state.keys[i] = migrationIndexCatalogKey{ - Position: i + 1, - IsExpression: isExpression, - Definition: definition, - TypeSchema: expected.resultType.schema, - TypeName: expected.resultType.name, - OpClassSchema: expected.operatorClass.schema, - OpClassName: expected.operatorClass.name, + Position: i + 1, + IsExpression: isExpression, + Definition: definition, + TypeSchema: expected.resultType.schema, + TypeName: expected.resultType.name, + OpClassSchema: expected.operatorClass.schema, + OpClassName: expected.operatorClass.name, + CollationSchema: expected.collation.schema, + CollationName: expected.collation.name, } } if requirement.predicateCanonical != "" { state.predicate = sql.NullString{ - String: "((reason)::text = ANY ((ARRAY['account_share_mode_seat_prepay'::character varying, 'account_share_mode_seat_refund'::character varying, 'account_share_mode_seat_waiver_refund'::character varying])::text[]))", + String: requirement.predicateCanonical, Valid: true, } - if requirement.index.name == "idx_user_balance_ledger_seat_membership_created_at" { - state.predicate.String += " AND (NULLIF((metadata ->> 'membership_id'::text), ''::text) IS NOT NULL)" - } } if mutate != nil { mutate(&state) @@ -217,6 +221,22 @@ func TestCanonicalizeMigrationIndexExpressionNormalizesPostgreSQLVarcharCasts(t require.NotEqual(t, canonicalizeMigrationIndexExpression("reason50"), canonicalizeMigrationIndexExpression("reason::varchar(50)")) require.NotEqual(t, canonicalizeMigrationIndexExpression("reason50"), canonicalizeMigrationIndexExpression("reason::character varying(50)")) }) + + t.Run("OpenAI身份索引表达式兼容PostgreSQL规范化输出", func(t *testing.T) { + actualOrganization := "lower(NULLIF(btrim((credentials ->> 'organization_id'::text)), ''::text))" + require.Equal( + t, + canonicalizeMigrationIndexExpression(openAIOwnedOrganizationExpression), + canonicalizeMigrationIndexExpression(actualOrganization), + ) + + actualAgentPredicate := "(deleted_at IS NULL) AND (owner_user_id IS NOT NULL) AND ((platform)::text = 'openai'::text) AND ((type)::text = 'oauth'::text) AND (lower(NULLIF(btrim((credentials ->> 'auth_mode'::text)), ''::text)) = 'agentidentity'::text) AND (NULLIF(btrim((credentials ->> 'chatgpt_account_id'::text)), ''::text) IS NOT NULL)" + require.Equal( + t, + canonicalizeMigrationIndexExpression(openAIOwnedAgentIdentityPredicate), + canonicalizeMigrationIndexExpression(actualAgentPredicate), + ) + }) } func TestPrepareAccountShareSeatCostIndexesMigrationDropsInvalidIndexBeforeRetry(t *testing.T) { @@ -627,6 +647,186 @@ DROP INDEX CONCURRENTLY IF EXISTS idx_accounts_owned_openai_chatgpt_account_id_u require.NoError(t, mock.ExpectationsWereMet()) } +func TestPrepareOpenAIOwnedAgentIdentityUniqueMigrationFailsFastOnDuplicates(t *testing.T) { + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { _ = db.Close() }() + + mock.ExpectQuery(openAIAgentIdentityDuplicatePrecheckPattern). + WillReturnRows(sqlmock.NewRows([]string{"identity_name", "owner_user_id", "duplicate_count", "sample_ids"}). + AddRow("openai.agent_identity_team", int64(101), 2, "41,42")) + + err = prepareNonTransactionalMigration(context.Background(), db, openAIOwnedAgentIdentityUniqueMigration) + require.Error(t, err) + require.Contains(t, err.Error(), "duplicate OpenAI owned Agent Identity migration identities") + require.Contains(t, err.Error(), "openai.agent_identity_team") + require.Contains(t, err.Error(), "sample_account_ids=41,42") + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestPrepareOpenAIOwnedAgentIdentityUniqueMigrationDropsInvalidIndexesBeforeRetry(t *testing.T) { + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { _ = db.Close() }() + + mock.ExpectQuery(openAIAgentIdentityDuplicatePrecheckPattern). + WillReturnRows(sqlmock.NewRows([]string{"identity_name", "owner_user_id", "duplicate_count", "sample_ids"})) + for i, requirement := range openAIOwnedAgentIdentityUniqueIndexRequirements { + invalid := i == 0 + mock.ExpectQuery("SELECT EXISTS \\("). + WithArgs(requirement.index.name). + WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(invalid)) + if i == 0 { + mock.ExpectExec(regexp.QuoteMeta("DROP INDEX CONCURRENTLY IF EXISTS " + quoteMigrationCatalogName(requirement.index))). + WillReturnResult(sqlmock.NewResult(0, 0)) + continue + } + mock.ExpectQuery("SELECT\\s+tbl_ns\\.nspname"). + WithArgs(requirement.index.schema, requirement.index.name). + WillReturnError(sql.ErrNoRows) + } + + err = prepareNonTransactionalMigration(context.Background(), db, openAIOwnedAgentIdentityUniqueMigration) + require.NoError(t, err) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestPrepareOpenAIOwnedAgentIdentityUniqueMigrationRejectsValidWrongDefinition(t *testing.T) { + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { _ = db.Close() }() + + mock.ExpectQuery(openAIAgentIdentityDuplicatePrecheckPattern). + WillReturnRows(sqlmock.NewRows([]string{"identity_name", "owner_user_id", "duplicate_count", "sample_ids"})) + requirement := openAIOwnedAgentIdentityUniqueIndexRequirements[0] + mock.ExpectQuery("SELECT EXISTS \\("). + WithArgs(requirement.index.name). + WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(false)) + expectMigrationIndexCatalogQuery(mock, requirement, matchingMigrationIndexCatalogRows(t, requirement, func(state *migrationIndexCatalogState) { + state.keys[1].Definition = "lower(credentials->>'wrong_field')" + })) + + err = prepareNonTransactionalMigration(context.Background(), db, openAIOwnedAgentIdentityUniqueMigration) + require.ErrorContains(t, err, "valid but unexpected definition") + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestVerifyOpenAIOwnedAgentIdentityUniqueMigrationRequiresExactDefinitions(t *testing.T) { + t.Run("all indexes match", func(t *testing.T) { + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { _ = db.Close() }() + + for _, requirement := range openAIOwnedAgentIdentityUniqueIndexRequirements { + expectMigrationIndexCatalogQuery(mock, requirement, matchingMigrationIndexCatalogRows(t, requirement, nil)) + } + + err = verifyNonTransactionalMigrationResult(context.Background(), db, openAIOwnedAgentIdentityUniqueMigration) + require.NoError(t, err) + require.NoError(t, mock.ExpectationsWereMet()) + }) + + t.Run("missing index fails verification", func(t *testing.T) { + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { _ = db.Close() }() + + requirement := openAIOwnedAgentIdentityUniqueIndexRequirements[0] + mock.ExpectQuery("SELECT\\s+tbl_ns\\.nspname"). + WithArgs(requirement.index.schema, requirement.index.name). + WillReturnError(sql.ErrNoRows) + + err = verifyNonTransactionalMigrationResult(context.Background(), db, openAIOwnedAgentIdentityUniqueMigration) + require.ErrorContains(t, err, "missing, invalid, or do not match migration 217") + require.NoError(t, mock.ExpectationsWereMet()) + }) + + t.Run("wrong expression fails verification", func(t *testing.T) { + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { _ = db.Close() }() + + requirement := openAIOwnedAgentIdentityUniqueIndexRequirements[0] + expectMigrationIndexCatalogQuery(mock, requirement, matchingMigrationIndexCatalogRows(t, requirement, func(state *migrationIndexCatalogState) { + state.keys[1].Definition = "lower(credentials->>'wrong_field')" + })) + + err = verifyNonTransactionalMigrationResult(context.Background(), db, openAIOwnedAgentIdentityUniqueMigration) + require.ErrorContains(t, err, "do not match migration 217") + require.NoError(t, mock.ExpectationsWereMet()) + }) +} + +func TestOpenAIOwnedAgentIdentityIndexMatcherRejectsCatalogDrift(t *testing.T) { + requirement := openAIOwnedAgentIdentityUniqueIndexRequirements[0] + tests := []struct { + name string + mutate func(*migrationIndexCatalogState) + }{ + {name: "not unique", mutate: func(state *migrationIndexCatalogState) { state.unique = false }}, + {name: "not ready", mutate: func(state *migrationIndexCatalogState) { state.ready = false }}, + {name: "not valid", mutate: func(state *migrationIndexCatalogState) { state.valid = false }}, + {name: "wrong result type", mutate: func(state *migrationIndexCatalogState) { state.keys[1].TypeName = "varchar" }}, + {name: "wrong operator class", mutate: func(state *migrationIndexCatalogState) { state.keys[1].OpClassName = "varchar_ops" }}, + {name: "wrong collation", mutate: func(state *migrationIndexCatalogState) { state.keys[1].CollationName = "C" }}, + {name: "wrong predicate", mutate: func(state *migrationIndexCatalogState) { state.predicate.String = "owner_user_id IS NULL" }}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { _ = db.Close() }() + + expectMigrationIndexCatalogQuery(mock, requirement, matchingMigrationIndexCatalogRows(t, requirement, test.mutate)) + state, err := loadMigrationIndexCatalogState(context.Background(), db, requirement.index) + require.NoError(t, err) + require.False(t, migrationIndexMatchesRequirement(state, requirement)) + require.NoError(t, mock.ExpectationsWereMet()) + }) + } +} + +func TestApplyOpenAIOwnedAgentIdentityMigrationVerifiesNewIndexesBeforeProtectedDrops(t *testing.T) { + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { _ = db.Close() }() + + prepareMigrationsBootstrapExpectations(mock) + mock.ExpectQuery("SELECT checksum FROM schema_migrations WHERE filename = \\$1"). + WithArgs(openAIOwnedAgentIdentityUniqueMigration). + WillReturnError(sql.ErrNoRows) + mock.ExpectQuery(openAIAgentIdentityDuplicatePrecheckPattern). + WillReturnRows(sqlmock.NewRows([]string{"identity_name", "owner_user_id", "duplicate_count", "sample_ids"})) + for _, requirement := range openAIOwnedAgentIdentityUniqueIndexRequirements { + mock.ExpectQuery("SELECT EXISTS \\("). + WithArgs(requirement.index.name). + WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(false)) + mock.ExpectQuery("SELECT\\s+tbl_ns\\.nspname"). + WithArgs(requirement.index.schema, requirement.index.name). + WillReturnError(sql.ErrNoRows) + } + mock.ExpectExec("CREATE UNIQUE INDEX CONCURRENTLY IF NOT EXISTS idx_test"). + WillReturnResult(sqlmock.NewResult(0, 0)) + firstRequirement := openAIOwnedAgentIdentityUniqueIndexRequirements[0] + mock.ExpectQuery("SELECT\\s+tbl_ns\\.nspname"). + WithArgs(firstRequirement.index.schema, firstRequirement.index.name). + WillReturnError(sql.ErrNoRows) + mock.ExpectExec("SELECT pg_advisory_unlock\\(\\$1\\)"). + WithArgs(migrationsAdvisoryLockID). + WillReturnResult(sqlmock.NewResult(0, 1)) + + fsys := fstest.MapFS{ + openAIOwnedAgentIdentityUniqueMigration: &fstest.MapFile{ + Data: []byte("CREATE UNIQUE INDEX CONCURRENTLY IF NOT EXISTS idx_test ON accounts (owner_user_id);\nDROP INDEX CONCURRENTLY IF EXISTS idx_protected_old;\n"), + }, + } + err = applyMigrationsFS(context.Background(), db, fsys) + + require.ErrorContains(t, err, "before dropping protected indexes") + require.NoError(t, mock.ExpectationsWereMet()) +} + func TestApplyMigrationsFS_TransactionalMigration(t *testing.T) { db, mock, err := sqlmock.New() require.NoError(t, err) @@ -658,6 +858,150 @@ func TestApplyMigrationsFS_TransactionalMigration(t *testing.T) { require.NoError(t, mock.ExpectationsWereMet()) } +func TestApplyMigrationsFSUsageLogImageInputTokensPinsPublicSchemaAndValidatesColumns(t *testing.T) { + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { _ = db.Close() }() + migrationSQL, err := migrations.FS.ReadFile(usageLogImageInputTokensMigration) + require.NoError(t, err) + require.Contains(t, string(migrationSQL), "ADD COLUMN IF NOT EXISTS image_input_tokens") + require.Contains(t, string(migrationSQL), "ADD COLUMN IF NOT EXISTS image_input_cost") + + prepareMigrationsBootstrapExpectations(mock) + mock.ExpectQuery("SELECT checksum FROM schema_migrations WHERE filename = \\$1"). + WithArgs(usageLogImageInputTokensMigration). + WillReturnError(sql.ErrNoRows) + mock.ExpectBegin() + mock.ExpectExec("SET LOCAL search_path = pg_catalog, public, pg_temp"). + WillReturnResult(sqlmock.NewResult(0, 0)) + mock.ExpectExec("ALTER TABLE usage_logs ADD COLUMN IF NOT EXISTS image_input_tokens"). + WillReturnResult(sqlmock.NewResult(0, 0)) + mock.ExpectQuery("SELECT\\s+a\\.attname"). + WillReturnRows(sqlmock.NewRows([]string{"attname", "format_type", "attnotnull", "default"}). + AddRow("image_input_cost", "numeric(20,10)", true, "0"). + AddRow("image_input_tokens", "integer", true, "0")) + mock.ExpectExec("INSERT INTO schema_migrations \\(filename, checksum\\) VALUES \\(\\$1, \\$2\\)"). + WithArgs(usageLogImageInputTokensMigration, sqlmock.AnyArg()). + WillReturnResult(sqlmock.NewResult(1, 1)) + mock.ExpectCommit() + mock.ExpectExec("SELECT pg_advisory_unlock\\(\\$1\\)"). + WithArgs(migrationsAdvisoryLockID). + WillReturnResult(sqlmock.NewResult(0, 1)) + + fsys := fstest.MapFS{ + usageLogImageInputTokensMigration: &fstest.MapFile{ + Data: migrationSQL, + }, + } + + err = applyMigrationsFS(context.Background(), db, fsys) + require.NoError(t, err) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestVerifyTransactionalMigrationResultRejectsMalformedUsageLogImageInputColumns(t *testing.T) { + tests := []struct { + name string + costType string + costNotNull bool + costDefault string + includeCostCol bool + wantErr string + }{ + { + name: "wrong default", + costType: "numeric(20,10)", + costNotNull: true, + costDefault: "1", + includeCostCol: true, + wantErr: "public.usage_logs.image_input_cost has unexpected definition", + }, + { + name: "wrong type", + costType: "numeric(20,9)", + costNotNull: true, + costDefault: "0", + includeCostCol: true, + wantErr: "public.usage_logs.image_input_cost has unexpected definition", + }, + { + name: "nullable", + costType: "numeric(20,10)", + costNotNull: false, + costDefault: "0", + includeCostCol: true, + wantErr: "public.usage_logs.image_input_cost has unexpected definition", + }, + { + name: "missing default", + costType: "numeric(20,10)", + costNotNull: true, + costDefault: "", + includeCostCol: true, + wantErr: "public.usage_logs.image_input_cost has unexpected definition", + }, + { + name: "missing column", + includeCostCol: false, + wantErr: "public.usage_logs.image_input_cost is missing", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { _ = db.Close() }() + + rows := sqlmock.NewRows([]string{"attname", "format_type", "attnotnull", "default"}). + AddRow("image_input_tokens", "integer", true, "0") + if test.includeCostCol { + rows.AddRow("image_input_cost", test.costType, test.costNotNull, test.costDefault) + } + mock.ExpectQuery("SELECT\\s+a\\.attname").WillReturnRows(rows) + + err = verifyTransactionalMigrationResult(context.Background(), db, usageLogImageInputTokensMigration) + require.ErrorContains(t, err, test.wantErr) + require.NoError(t, mock.ExpectationsWereMet()) + }) + } +} + +func TestApplyMigrationsFSUsageLogImageInputTokensRollsBackWhenValidationFails(t *testing.T) { + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { _ = db.Close() }() + migrationSQL, err := migrations.FS.ReadFile(usageLogImageInputTokensMigration) + require.NoError(t, err) + + prepareMigrationsBootstrapExpectations(mock) + mock.ExpectQuery("SELECT checksum FROM schema_migrations WHERE filename = \\$1"). + WithArgs(usageLogImageInputTokensMigration). + WillReturnError(sql.ErrNoRows) + mock.ExpectBegin() + mock.ExpectExec("SET LOCAL search_path = pg_catalog, public, pg_temp"). + WillReturnResult(sqlmock.NewResult(0, 0)) + mock.ExpectExec("ALTER TABLE usage_logs ADD COLUMN IF NOT EXISTS image_input_tokens"). + WillReturnResult(sqlmock.NewResult(0, 0)) + mock.ExpectQuery("SELECT\\s+a\\.attname"). + WillReturnRows(sqlmock.NewRows([]string{"attname", "format_type", "attnotnull", "default"}). + AddRow("image_input_cost", "numeric(20,10)", false, "0"). + AddRow("image_input_tokens", "integer", true, "0")) + mock.ExpectRollback() + mock.ExpectExec("SELECT pg_advisory_unlock\\(\\$1\\)"). + WithArgs(migrationsAdvisoryLockID). + WillReturnResult(sqlmock.NewResult(0, 1)) + + fsys := fstest.MapFS{ + usageLogImageInputTokensMigration: &fstest.MapFile{Data: migrationSQL}, + } + err = applyMigrationsFS(context.Background(), db, fsys) + + require.ErrorContains(t, err, "verify migration "+usageLogImageInputTokensMigration) + require.ErrorContains(t, err, "public.usage_logs.image_input_cost has unexpected definition") + require.NoError(t, mock.ExpectationsWereMet()) +} + func prepareMigrationsBootstrapExpectations(mock sqlmock.Sqlmock) { mock.ExpectQuery("SELECT pg_try_advisory_lock\\(\\$1\\)"). WithArgs(migrationsAdvisoryLockID). diff --git a/backend/internal/repository/scheduler_cache.go b/backend/internal/repository/scheduler_cache.go index d1a847fbc..c5a15a1ee 100644 --- a/backend/internal/repository/scheduler_cache.go +++ b/backend/internal/repository/scheduler_cache.go @@ -1118,6 +1118,9 @@ func buildSchedulerMetadataAccount(account service.Account) service.Account { Name: account.Name, Platform: account.Platform, Type: account.Type, + OwnerUserID: account.OwnerUserID, + ShareMode: account.ShareMode, + ShareStatus: account.ShareStatus, Concurrency: account.Concurrency, LoadFactor: account.LoadFactor, Priority: account.Priority, @@ -1242,6 +1245,8 @@ func filterSchedulerExtra(extra map[string]any) map[string]any { "codex_7d_reset_at", "codex_7d_reset_after_seconds", "codex_7d_limit_percent", + service.GrokMediaEligibleExtraKey, + "grok_billing_snapshot", } filtered := make(map[string]any) for _, key := range keys { diff --git a/backend/internal/repository/scheduler_cache_unit_test.go b/backend/internal/repository/scheduler_cache_unit_test.go index 1f8e30fad..257b1fbc9 100644 --- a/backend/internal/repository/scheduler_cache_unit_test.go +++ b/backend/internal/repository/scheduler_cache_unit_test.go @@ -4,6 +4,7 @@ package repository import ( "context" + "encoding/json" "testing" "time" @@ -79,6 +80,30 @@ func TestBuildSchedulerMetadataAccount_KeepsSlimGroupMembership(t *testing.T) { require.Nil(t, got.Groups) } +func TestBuildSchedulerMetadataAccount_KeepsOwnedShareVisibilityFields(t *testing.T) { + ownerUserID := int64(73) + account := service.Account{ + ID: 42, + Platform: service.PlatformOpenAI, + Type: service.AccountTypeOAuth, + OwnerUserID: &ownerUserID, + ShareMode: service.AccountShareModePublic, + ShareStatus: service.AccountShareStatusPending, + } + + metadata := buildSchedulerMetadataAccount(account) + payload, err := json.Marshal(metadata) + require.NoError(t, err) + + var decoded service.Account + require.NoError(t, json.Unmarshal(payload, &decoded)) + require.NotNil(t, decoded.OwnerUserID) + require.Equal(t, ownerUserID, *decoded.OwnerUserID) + require.Equal(t, service.AccountShareModePublic, decoded.ShareMode) + require.Equal(t, service.AccountShareStatusPending, decoded.ShareStatus) + require.False(t, decoded.IsVisibleToConsumer(ownerUserID+1)) +} + func TestSchedulerCacheWriteAccountsSkipsUnencodableTime(t *testing.T) { rdb := redis.NewClient(&redis.Options{Addr: "127.0.0.1:0"}) t.Cleanup(func() { _ = rdb.Close() }) diff --git a/backend/internal/repository/usage_billing_repo.go b/backend/internal/repository/usage_billing_repo.go index 35dcd5dfc..e58d895b5 100644 --- a/backend/internal/repository/usage_billing_repo.go +++ b/backend/internal/repository/usage_billing_repo.go @@ -456,6 +456,8 @@ func usageBillingUsageLogInsertQuery() string { cache_creation_1h_tokens, image_output_tokens, image_output_cost, + image_input_tokens, + image_input_cost, input_cost, output_cost, cache_creation_cost, @@ -495,7 +497,7 @@ func usageBillingUsageLogInsertQuery() string { $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, - $24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50 + $24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50, $51, $52 ) ON CONFLICT (request_id, api_key_id) DO NOTHING RETURNING id, created_at diff --git a/backend/internal/repository/usage_log_repo.go b/backend/internal/repository/usage_log_repo.go index 9668a4764..ca1e3b42f 100644 --- a/backend/internal/repository/usage_log_repo.go +++ b/backend/internal/repository/usage_log_repo.go @@ -30,7 +30,7 @@ import ( gocache "github.com/patrickmn/go-cache" ) -const usageLogSelectColumns = "id, user_id, api_key_id, account_id, request_id, model, requested_model, upstream_model, group_id, subscription_id, input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, cache_creation_5m_tokens, cache_creation_1h_tokens, image_output_tokens, image_output_cost, input_cost, output_cost, cache_creation_cost, cache_read_cost, total_cost, actual_cost, rate_multiplier, rate_multiplier_source, account_rate_multiplier, billing_type, request_type, stream, openai_ws_mode, duration_ms, first_token_ms, user_agent, ip_address, image_count, image_size, video_count, video_resolution, video_duration_seconds, service_tier, reasoning_effort, inbound_endpoint, upstream_endpoint, cache_ttl_overridden, channel_id, model_mapping_chain, billing_tier, billing_mode, account_stats_cost, created_at" +const usageLogSelectColumns = "id, user_id, api_key_id, account_id, request_id, model, requested_model, upstream_model, group_id, subscription_id, input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, cache_creation_5m_tokens, cache_creation_1h_tokens, image_output_tokens, image_output_cost, image_input_tokens, image_input_cost, input_cost, output_cost, cache_creation_cost, cache_read_cost, total_cost, actual_cost, rate_multiplier, rate_multiplier_source, account_rate_multiplier, billing_type, request_type, stream, openai_ws_mode, duration_ms, first_token_ms, user_agent, ip_address, image_count, image_size, video_count, video_resolution, video_duration_seconds, service_tier, reasoning_effort, inbound_endpoint, upstream_endpoint, cache_ttl_overridden, channel_id, model_mapping_chain, billing_tier, billing_mode, account_stats_cost, created_at" // usageLogInsertArgTypes must stay in the same order as: // 1. prepareUsageLogInsert().args @@ -57,6 +57,8 @@ var usageLogInsertArgTypes = [...]string{ "integer", // cache_creation_1h_tokens "integer", // image_output_tokens "numeric", // image_output_cost + "integer", // image_input_tokens + "numeric", // image_input_cost "numeric", // input_cost "numeric", // output_cost "numeric", // cache_creation_cost @@ -116,6 +118,35 @@ func safeDateFormat(granularity string) string { return "YYYY-MM-DD" } +// dashboardBucketKeyExpression returns a trusted SQL expression for dashboard trend grouping. +// timestampColumn must be a static column expression owned by this package, never request input. +func dashboardBucketKeyExpression(granularity, timestampColumn string) string { + if granularity == "hour" { + // Keep repeated wall-clock hours distinct during a DST rollback by grouping on + // the absolute UTC hour. The client timezone is applied only to the label. + return fmt.Sprintf("DATE_TRUNC('hour', %s AT TIME ZONE 'UTC')", timestampColumn) + } + + unit := "day" + switch granularity { + case "week": + unit = "week" + case "month": + unit = "month" + } + return fmt.Sprintf("DATE_TRUNC('%s', %s AT TIME ZONE $4)", unit, timestampColumn) +} + +func dashboardBucketLabelExpression(granularity, bucketColumn string) string { + if granularity == "hour" { + return fmt.Sprintf( + "TO_CHAR((%s AT TIME ZONE 'UTC') AT TIME ZONE $4, 'YYYY-MM-DD HH24:MI')", + bucketColumn, + ) + } + return fmt.Sprintf("TO_CHAR(%s, '%s')", bucketColumn, safeDateFormat(granularity)) +} + // appendRawUsageLogModelWhereCondition keeps direct model filters on the raw model column for backward // compatibility with historical rows. Requested/upstream analytics must use // resolveModelDimensionExpression instead. @@ -342,6 +373,8 @@ func (r *usageLogRepository) createSingle(ctx context.Context, sqlq sqlExecutor, cache_creation_1h_tokens, image_output_tokens, image_output_cost, + image_input_tokens, + image_input_cost, input_cost, output_cost, cache_creation_cost, @@ -381,7 +414,7 @@ func (r *usageLogRepository) createSingle(ctx context.Context, sqlq sqlExecutor, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, - $24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50 + $24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50, $51, $52 ) ON CONFLICT (request_id, api_key_id) DO NOTHING RETURNING id, created_at @@ -784,6 +817,8 @@ func buildUsageLogBatchInsertQuery(keys []string, preparedByKey map[string]usage cache_creation_1h_tokens, image_output_tokens, image_output_cost, + image_input_tokens, + image_input_cost, input_cost, output_cost, cache_creation_cost, @@ -819,7 +854,7 @@ func buildUsageLogBatchInsertQuery(keys []string, preparedByKey map[string]usage created_at ) AS (VALUES `) - args := make([]any, 0, len(keys)*47) + args := make([]any, 0, len(keys)*53) argPos := 1 for idx, key := range keys { if idx > 0 { @@ -865,6 +900,8 @@ func buildUsageLogBatchInsertQuery(keys []string, preparedByKey map[string]usage cache_creation_1h_tokens, image_output_tokens, image_output_cost, + image_input_tokens, + image_input_cost, input_cost, output_cost, cache_creation_cost, @@ -917,6 +954,8 @@ func buildUsageLogBatchInsertQuery(keys []string, preparedByKey map[string]usage cache_creation_1h_tokens, image_output_tokens, image_output_cost, + image_input_tokens, + image_input_cost, input_cost, output_cost, cache_creation_cost, @@ -1009,6 +1048,8 @@ func buildUsageLogBestEffortInsertQuery(preparedList []usageLogInsertPrepared) ( cache_creation_1h_tokens, image_output_tokens, image_output_cost, + image_input_tokens, + image_input_cost, input_cost, output_cost, cache_creation_cost, @@ -1044,7 +1085,7 @@ func buildUsageLogBestEffortInsertQuery(preparedList []usageLogInsertPrepared) ( created_at ) AS (VALUES `) - args := make([]any, 0, len(preparedList)*47) + args := make([]any, 0, len(preparedList)*52) argPos := 1 for idx, prepared := range preparedList { if idx > 0 { @@ -1087,6 +1128,8 @@ func buildUsageLogBestEffortInsertQuery(preparedList []usageLogInsertPrepared) ( cache_creation_1h_tokens, image_output_tokens, image_output_cost, + image_input_tokens, + image_input_cost, input_cost, output_cost, cache_creation_cost, @@ -1139,6 +1182,8 @@ func buildUsageLogBestEffortInsertQuery(preparedList []usageLogInsertPrepared) ( cache_creation_1h_tokens, image_output_tokens, image_output_cost, + image_input_tokens, + image_input_cost, input_cost, output_cost, cache_creation_cost, @@ -1199,6 +1244,8 @@ func execUsageLogInsertNoResult(ctx context.Context, sqlq sqlExecutor, prepared cache_creation_1h_tokens, image_output_tokens, image_output_cost, + image_input_tokens, + image_input_cost, input_cost, output_cost, cache_creation_cost, @@ -1238,7 +1285,7 @@ func execUsageLogInsertNoResult(ctx context.Context, sqlq sqlExecutor, prepared $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, - $24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50 + $24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50, $51, $52 ) ON CONFLICT (request_id, api_key_id) DO NOTHING `, prepared.args...) @@ -1310,6 +1357,8 @@ func prepareUsageLogInsert(log *service.UsageLog) usageLogInsertPrepared { log.CacheCreation1hTokens, log.ImageOutputTokens, log.ImageOutputCost, + log.ImageInputTokens, + log.ImageInputCost, log.InputCost, log.OutputCost, log.CacheCreationCost, @@ -3020,27 +3069,44 @@ func (r *usageLogRepository) GetAPIKeyDashboardStats(ctx context.Context, apiKey } // GetUserUsageTrendByUserID 获取指定用户的使用趋势 -func (r *usageLogRepository) GetUserUsageTrendByUserID(ctx context.Context, userID int64, startTime, endTime time.Time, granularity string) (results []TrendDataPoint, err error) { - dateFormat := safeDateFormat(granularity) +func (r *usageLogRepository) GetUserUsageTrendByUserID(ctx context.Context, userID int64, startTime, endTime time.Time, granularity string, location *time.Location) (results []TrendDataPoint, err error) { + if location == nil { + return nil, fmt.Errorf("dashboard timezone location is required") + } + bucketKeyExpression := dashboardBucketKeyExpression(granularity, "created_at") + bucketLabelExpression := dashboardBucketLabelExpression(granularity, "bucket_key") query := fmt.Sprintf(` + WITH bucketed AS ( + SELECT + %s AS bucket_key, + COUNT(*) AS requests, + COALESCE(SUM(input_tokens), 0) AS input_tokens, + COALESCE(SUM(output_tokens), 0) AS output_tokens, + COALESCE(SUM(cache_creation_tokens), 0) AS cache_creation_tokens, + COALESCE(SUM(cache_read_tokens), 0) AS cache_read_tokens, + COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_tokens, + COALESCE(SUM(total_cost), 0) AS cost, + COALESCE(SUM(actual_cost), 0) AS actual_cost + FROM usage_logs + WHERE user_id = $1 AND created_at >= $2 AND created_at < $3 + GROUP BY bucket_key + ) SELECT - TO_CHAR(created_at, '%s') as date, - COUNT(*) as requests, - COALESCE(SUM(input_tokens), 0) as input_tokens, - COALESCE(SUM(output_tokens), 0) as output_tokens, - COALESCE(SUM(cache_creation_tokens), 0) as cache_creation_tokens, - COALESCE(SUM(cache_read_tokens), 0) as cache_read_tokens, - COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) as total_tokens, - COALESCE(SUM(total_cost), 0) as cost, - COALESCE(SUM(actual_cost), 0) as actual_cost - FROM usage_logs - WHERE user_id = $1 AND created_at >= $2 AND created_at < $3 - GROUP BY date - ORDER BY date ASC - `, dateFormat) - - rows, err := r.sql.QueryContext(ctx, query, userID, startTime, endTime) + %s AS date, + requests, + input_tokens, + output_tokens, + cache_creation_tokens, + cache_read_tokens, + total_tokens, + cost, + actual_cost + FROM bucketed + ORDER BY bucket_key ASC + `, bucketKeyExpression, bucketLabelExpression) + + rows, err := r.sql.QueryContext(ctx, query, userID, startTime, endTime, location.String()) if err != nil { return nil, err } @@ -3101,10 +3167,13 @@ func (r *usageLogRepository) GetUserModelStats(ctx context.Context, userID int64 } // GetUserAccountSharingDashboard returns owned-account self usage and external public-share settlement stats. -func (r *usageLogRepository) GetUserAccountSharingDashboard(ctx context.Context, userID int64, startTime, endTime time.Time, granularity string) (*usagestats.AccountSharingDashboardStats, error) { +func (r *usageLogRepository) GetUserAccountSharingDashboard(ctx context.Context, userID int64, startTime, endTime time.Time, granularity string, location *time.Location) (*usagestats.AccountSharingDashboardStats, error) { if userID <= 0 { return nil, fmt.Errorf("user id must be positive") } + if location == nil { + return nil, fmt.Errorf("dashboard timezone location is required") + } if startTime.IsZero() { startTime = time.Now().AddDate(0, 0, -7) } @@ -3116,7 +3185,7 @@ func (r *usageLogRepository) GetUserAccountSharingDashboard(ctx context.Context, if err != nil { return nil, err } - trend, err := r.getUserAccountSharingTrend(ctx, userID, startTime, endTime, granularity) + trend, err := r.getUserAccountSharingTrend(ctx, userID, startTime, endTime, granularity, location) if err != nil { return nil, err } @@ -3129,8 +3198,8 @@ func (r *usageLogRepository) GetUserAccountSharingDashboard(ctx context.Context, Summary: summary, Accounts: accounts, Trend: trend, - StartDate: startTime.Format("2006-01-02"), - EndDate: endDisplay.Format("2006-01-02"), + StartDate: startTime.In(location).Format("2006-01-02"), + EndDate: endDisplay.In(location).Format("2006-01-02"), Granularity: granularity, }, nil } @@ -3252,12 +3321,14 @@ func (r *usageLogRepository) getUserAccountSharingAccountStats(ctx context.Conte return accounts, summary, nil } -func (r *usageLogRepository) getUserAccountSharingTrend(ctx context.Context, userID int64, startTime, endTime time.Time, granularity string) (results []usagestats.AccountSharingTrendPoint, err error) { - dateFormat := safeDateFormat(granularity) +func (r *usageLogRepository) getUserAccountSharingTrend(ctx context.Context, userID int64, startTime, endTime time.Time, granularity string, location *time.Location) (results []usagestats.AccountSharingTrendPoint, err error) { + bucketKeySelf := dashboardBucketKeyExpression(granularity, "ul.created_at") + bucketKeyExternal := dashboardBucketKeyExpression(granularity, "created_at") + bucketLabel := dashboardBucketLabelExpression(granularity, "COALESCE(s.bucket_key, e.bucket_key)") query := fmt.Sprintf(` WITH self_usage AS ( SELECT - TO_CHAR(ul.created_at, '%s') AS date, + %s AS bucket_key, COUNT(*) AS self_requests, COALESCE(SUM(ul.input_tokens + ul.output_tokens + ul.cache_creation_tokens + ul.cache_read_tokens), 0) AS self_tokens, COALESCE(SUM(ul.actual_cost), 0) AS self_actual_cost, @@ -3268,11 +3339,11 @@ func (r *usageLogRepository) getUserAccountSharingTrend(ctx context.Context, use AND ul.user_id = $1 AND ul.created_at >= $2 AND ul.created_at < $3 - GROUP BY date + GROUP BY bucket_key ), external_usage AS ( SELECT - TO_CHAR(created_at, '%s') AS date, + %s AS bucket_key, COUNT(*) AS external_requests, COALESCE(SUM(consumer_charge), 0) AS external_consumer_charge, COALESCE(SUM(account_cost), 0) AS external_account_cost, @@ -3284,10 +3355,10 @@ func (r *usageLogRepository) getUserAccountSharingTrend(ctx context.Context, use AND status = 'applied' AND created_at >= $2 AND created_at < $3 - GROUP BY date + GROUP BY bucket_key ) SELECT - COALESCE(s.date, e.date) AS date, + %s AS date, COALESCE(s.self_requests, 0), COALESCE(s.self_tokens, 0), COALESCE(s.self_actual_cost, 0), @@ -3298,11 +3369,11 @@ func (r *usageLogRepository) getUserAccountSharingTrend(ctx context.Context, use COALESCE(e.external_owner_credit, 0), COALESCE(e.external_platform_fee, 0) FROM self_usage s - FULL OUTER JOIN external_usage e ON e.date = s.date - ORDER BY date ASC - `, dateFormat, dateFormat) + FULL OUTER JOIN external_usage e ON e.bucket_key = s.bucket_key + ORDER BY COALESCE(s.bucket_key, e.bucket_key) ASC + `, bucketKeySelf, bucketKeyExternal, bucketLabel) - rows, err := r.sql.QueryContext(ctx, query, userID, startTime, endTime) + rows, err := r.sql.QueryContext(ctx, query, userID, startTime, endTime, location.String()) if err != nil { return nil, err } @@ -5833,6 +5904,8 @@ func scanUsageLog(scanner interface{ Scan(...any) error }) (*service.UsageLog, e cacheCreation1h int imageOutputTokens int imageOutputCost float64 + imageInputTokens int + imageInputCost float64 inputCost float64 outputCost float64 cacheCreationCost float64 @@ -5887,6 +5960,8 @@ func scanUsageLog(scanner interface{ Scan(...any) error }) (*service.UsageLog, e &cacheCreation1h, &imageOutputTokens, &imageOutputCost, + &imageInputTokens, + &imageInputCost, &inputCost, &outputCost, &cacheCreationCost, @@ -5939,6 +6014,8 @@ func scanUsageLog(scanner interface{ Scan(...any) error }) (*service.UsageLog, e CacheCreation1hTokens: cacheCreation1h, ImageOutputTokens: imageOutputTokens, ImageOutputCost: imageOutputCost, + ImageInputTokens: imageInputTokens, + ImageInputCost: imageInputCost, InputCost: inputCost, OutputCost: outputCost, CacheCreationCost: cacheCreationCost, diff --git a/backend/internal/repository/usage_log_repo_dashboard_timezone_test.go b/backend/internal/repository/usage_log_repo_dashboard_timezone_test.go new file mode 100644 index 000000000..edcc8db44 --- /dev/null +++ b/backend/internal/repository/usage_log_repo_dashboard_timezone_test.go @@ -0,0 +1,107 @@ +package repository + +import ( + "context" + "testing" + "time" + + "github.com/DATA-DOG/go-sqlmock" + "github.com/stretchr/testify/require" +) + +func TestDashboardBucketExpressionsKeepRollbackHoursDistinct(t *testing.T) { + hourKey := dashboardBucketKeyExpression("hour", "created_at") + hourLabel := dashboardBucketLabelExpression("hour", "bucket_key") + dayKey := dashboardBucketKeyExpression("day", "created_at") + + require.Equal(t, "DATE_TRUNC('hour', created_at AT TIME ZONE 'UTC')", hourKey) + require.Equal( + t, + "TO_CHAR((bucket_key AT TIME ZONE 'UTC') AT TIME ZONE $4, 'YYYY-MM-DD HH24:MI')", + hourLabel, + ) + require.Equal(t, "DATE_TRUNC('day', created_at AT TIME ZONE $4)", dayKey) + + location, err := time.LoadLocation("America/New_York") + require.NoError(t, err) + firstPhysicalHour := time.Date(2026, time.November, 1, 5, 0, 0, 0, time.UTC) + secondPhysicalHour := firstPhysicalHour.Add(time.Hour) + require.Equal( + t, + firstPhysicalHour.In(location).Format("2006-01-02 15:04"), + secondPhysicalHour.In(location).Format("2006-01-02 15:04"), + "the rollback hours intentionally share a display label", + ) + require.NotEqual(t, firstPhysicalHour, secondPhysicalHour, "their UTC bucket keys must remain distinct") +} + +func TestGetUserUsageTrendByUserIDUsesParameterizedTimezoneAndAbsoluteHourBucket(t *testing.T) { + db, mock := newSQLMock(t) + repo := &usageLogRepository{sql: db} + location, err := time.LoadLocation("America/New_York") + require.NoError(t, err) + start := time.Date(2026, time.November, 1, 4, 0, 0, 0, time.UTC) + end := start.Add(4 * time.Hour) + + mock.ExpectQuery(`(?s)DATE_TRUNC\('hour', created_at AT TIME ZONE 'UTC'\) AS bucket_key.*GROUP BY bucket_key.*TO_CHAR\(\(bucket_key AT TIME ZONE 'UTC'\) AT TIME ZONE \$4, 'YYYY-MM-DD HH24:MI'\) AS date.*ORDER BY bucket_key ASC`). + WithArgs(int64(42), start, end, "America/New_York"). + WillReturnRows(sqlmock.NewRows([]string{ + "date", + "requests", + "input_tokens", + "output_tokens", + "cache_creation_tokens", + "cache_read_tokens", + "total_tokens", + "cost", + "actual_cost", + })) + + trend, err := repo.GetUserUsageTrendByUserID( + context.Background(), + 42, + start, + end, + "hour", + location, + ) + require.NoError(t, err) + require.Empty(t, trend) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestGetUserAccountSharingTrendJoinsOnAbsoluteHourBucket(t *testing.T) { + db, mock := newSQLMock(t) + repo := &usageLogRepository{sql: db} + location, err := time.LoadLocation("America/New_York") + require.NoError(t, err) + start := time.Date(2026, time.November, 1, 4, 0, 0, 0, time.UTC) + end := start.Add(4 * time.Hour) + + mock.ExpectQuery(`(?s)DATE_TRUNC\('hour', ul\.created_at AT TIME ZONE 'UTC'\) AS bucket_key.*DATE_TRUNC\('hour', created_at AT TIME ZONE 'UTC'\) AS bucket_key.*FULL OUTER JOIN external_usage e ON e\.bucket_key = s\.bucket_key.*ORDER BY COALESCE\(s\.bucket_key, e\.bucket_key\) ASC`). + WithArgs(int64(42), start, end, "America/New_York"). + WillReturnRows(sqlmock.NewRows([]string{ + "date", + "self_requests", + "self_tokens", + "self_actual_cost", + "self_account_cost", + "external_requests", + "external_consumer_charge", + "external_account_cost", + "external_owner_credit", + "external_platform_fee", + })) + + trend, err := repo.getUserAccountSharingTrend( + context.Background(), + 42, + start, + end, + "hour", + location, + ) + require.NoError(t, err) + require.Empty(t, trend) + require.NoError(t, mock.ExpectationsWereMet()) +} diff --git a/backend/internal/repository/usage_log_repo_integration_test.go b/backend/internal/repository/usage_log_repo_integration_test.go index ee6f705da..a7f65afad 100644 --- a/backend/internal/repository/usage_log_repo_integration_test.go +++ b/backend/internal/repository/usage_log_repo_integration_test.go @@ -1283,7 +1283,7 @@ func (s *UsageLogRepoSuite) TestGetUserUsageTrendByUserID() { startTime := base.Add(-1 * time.Hour) endTime := base.Add(48 * time.Hour) - trend, err := s.repo.GetUserUsageTrendByUserID(s.ctx, user.ID, startTime, endTime, "day") + trend, err := s.repo.GetUserUsageTrendByUserID(s.ctx, user.ID, startTime, endTime, "day", time.UTC) s.Require().NoError(err, "GetUserUsageTrendByUserID") s.Require().Len(trend, 2) // 2 different days } @@ -1300,7 +1300,7 @@ func (s *UsageLogRepoSuite) TestGetUserUsageTrendByUserID_HourlyGranularity() { startTime := base.Add(-1 * time.Hour) endTime := base.Add(3 * time.Hour) - trend, err := s.repo.GetUserUsageTrendByUserID(s.ctx, user.ID, startTime, endTime, "hour") + trend, err := s.repo.GetUserUsageTrendByUserID(s.ctx, user.ID, startTime, endTime, "hour", time.UTC) s.Require().NoError(err, "GetUserUsageTrendByUserID hourly") s.Require().Len(trend, 3) // 3 different hours } diff --git a/backend/internal/repository/usage_log_repo_request_type_test.go b/backend/internal/repository/usage_log_repo_request_type_test.go index 2bfe33cfb..1a5341750 100644 --- a/backend/internal/repository/usage_log_repo_request_type_test.go +++ b/backend/internal/repository/usage_log_repo_request_type_test.go @@ -201,6 +201,8 @@ func TestUsageLogRepositoryCreateSyncRequestTypeAndLegacyFields(t *testing.T) { log.CacheCreation1hTokens, log.ImageOutputTokens, log.ImageOutputCost, + log.ImageInputTokens, + log.ImageInputCost, log.InputCost, log.OutputCost, log.CacheCreationCost, @@ -284,6 +286,8 @@ func TestUsageLogRepositoryCreate_PersistsServiceTier(t *testing.T) { log.CacheCreation1hTokens, log.ImageOutputTokens, log.ImageOutputCost, + log.ImageInputTokens, + log.ImageInputCost, log.InputCost, log.OutputCost, log.CacheCreationCost, @@ -381,6 +385,22 @@ func TestPrepareUsageLogInsert_ArgCountMatchesTypes(t *testing.T) { require.Len(t, prepared.args, len(usageLogInsertArgTypes)) } +func TestPrepareUsageLogInsert_PersistsImageInputUsage(t *testing.T) { + prepared := prepareUsageLogInsert(&service.UsageLog{ + UserID: 1, + APIKeyID: 2, + AccountID: 3, + RequestID: "req-image-input", + Model: "gpt-image-2", + ImageInputTokens: 352, + ImageInputCost: 0.002816, + CreatedAt: time.Date(2025, 1, 5, 12, 0, 0, 0, time.UTC), + }) + + require.Equal(t, 352, prepared.args[17]) + require.InDelta(t, 0.002816, prepared.args[18], 1e-15) +} + func TestUsageBillingUsageLogInsertQuery_ArgCountMatchesPreparedInsert(t *testing.T) { prepared := prepareUsageLogInsert(&service.UsageLog{ UserID: 1, @@ -844,6 +864,8 @@ func TestScanUsageLogRequestTypeAndLegacyFallback(t *testing.T) { 6, // cache_creation_1h_tokens 0, // image_output_tokens 0.0, // image_output_cost + 0, // image_input_tokens + 0.0, // image_input_cost 0.1, // input_cost 0.2, // output_cost 0.3, // cache_creation_cost @@ -901,6 +923,7 @@ func TestScanUsageLogRequestTypeAndLegacyFallback(t *testing.T) { sql.NullInt64{}, 1, 2, 3, 4, 5, 6, 0, 0.0, // image_output_tokens, image_output_cost + 0, 0.0, // image_input_tokens, image_input_cost 0.1, 0.2, 0.3, 0.4, 1.0, 0.9, 1.0, sql.NullString{}, // rate_multiplier_source @@ -953,6 +976,7 @@ func TestScanUsageLogRequestTypeAndLegacyFallback(t *testing.T) { sql.NullInt64{}, 1, 2, 3, 4, 5, 6, 0, 0.0, // image_output_tokens, image_output_cost + 0, 0.0, // image_input_tokens, image_input_cost 0.1, 0.2, 0.3, 0.4, 1.0, 0.9, 1.0, sql.NullString{}, // rate_multiplier_source diff --git a/backend/internal/repository/withdrawal_repo.go b/backend/internal/repository/withdrawal_repo.go index 90b889bb2..35afdb72d 100644 --- a/backend/internal/repository/withdrawal_repo.go +++ b/backend/internal/repository/withdrawal_repo.go @@ -45,7 +45,7 @@ SELECT EXISTS ( if pendingExists { return nil, service.ErrWithdrawalPendingExists } - if input.RateLimit.MaxRequests > 0 { + if input.RateLimit.MaxRequests > 0 && !input.RateLimit.ExemptsAmount(input.Amount) { if err := service.ValidateWithdrawalRateLimitConfig(input.RateLimit); err != nil { return nil, err } @@ -54,9 +54,11 @@ SELECT EXISTS ( SELECT COUNT(*) FROM user_withdrawal_requests WHERE user_id = $1 - AND created_at >= NOW() - ($2::integer * INTERVAL '1 day')`, + AND created_at >= NOW() - ($2::integer * INTERVAL '1 day') + AND ($3::numeric = 0 OR amount <= $3::numeric)`, input.UserID, input.RateLimit.WindowDays, + input.RateLimit.ExemptAmount, ).Scan(&recentRequestCount); err != nil { return nil, err } diff --git a/backend/internal/repository/withdrawal_repo_test.go b/backend/internal/repository/withdrawal_repo_test.go new file mode 100644 index 000000000..a6228bba5 --- /dev/null +++ b/backend/internal/repository/withdrawal_repo_test.go @@ -0,0 +1,82 @@ +package repository + +import ( + "context" + "errors" + "testing" + + sqlmock "github.com/DATA-DOG/go-sqlmock" + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/stretchr/testify/require" +) + +func TestWithdrawalSubmitRateLimitIncludesThresholdAmount(t *testing.T) { + repo, mock := newWithdrawalRateLimitRepository(t) + expectWithdrawalSubmitPreamble(mock, 1) + mock.ExpectQuery(`(?s)SELECT COUNT\(\*\).*amount <= \$3::numeric`). + WithArgs(int64(1), 7, 500.0). + WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(3)) + mock.ExpectRollback() + + result, err := repo.Submit(context.Background(), service.WithdrawalSubmitInput{ + UserID: 1, + Amount: 500, + PaymentMethod: "alipay", + RateLimit: service.WithdrawalRateLimitConfig{ + WindowDays: 7, + MaxRequests: 3, + ExemptAmount: 500, + }, + }) + + require.Nil(t, result) + require.Equal(t, "WITHDRAWAL_RATE_LIMIT_EXCEEDED", infraerrors.Reason(err)) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestWithdrawalSubmitSkipsRateLimitAboveThreshold(t *testing.T) { + repo, mock := newWithdrawalRateLimitRepository(t) + expectWithdrawalSubmitPreamble(mock, 1) + nextQueryErr := errors.New("reached post-rate-limit query") + mock.ExpectQuery(`(?s)SELECT EXISTS \(.*WHERE user_id = \$1.*\)`). + WithArgs(int64(1)). + WillReturnError(nextQueryErr) + mock.ExpectRollback() + + result, err := repo.Submit(context.Background(), service.WithdrawalSubmitInput{ + UserID: 1, + Amount: 500.01, + PaymentMethod: "alipay", + RateLimit: service.WithdrawalRateLimitConfig{ + WindowDays: 7, + MaxRequests: 3, + ExemptAmount: 500, + }, + }) + + require.Nil(t, result) + require.ErrorIs(t, err, nextQueryErr) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func newWithdrawalRateLimitRepository(t *testing.T) (*withdrawalRepository, sqlmock.Sqlmock) { + t.Helper() + db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp)) + require.NoError(t, err) + t.Cleanup(func() { + mock.ExpectClose() + require.NoError(t, db.Close()) + }) + return &withdrawalRepository{db: db}, mock +} + +func expectWithdrawalSubmitPreamble(mock sqlmock.Sqlmock, userID int64) { + mock.ExpectBegin() + mock.ExpectQuery(`(?s)SELECT email, balance::double precision.*FOR UPDATE`). + WithArgs(userID). + WillReturnRows(sqlmock.NewRows([]string{"email", "balance"}).AddRow("user@example.com", 1000.0)) + mock.ExpectQuery(`(?s)SELECT EXISTS \(.*status = \$2.*\)`). + WithArgs(userID, service.WithdrawalStatusPending). + WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(false)) +} diff --git a/backend/internal/server/api_contract_test.go b/backend/internal/server/api_contract_test.go index a8bec55ff..a44aab6b4 100644 --- a/backend/internal/server/api_contract_test.go +++ b/backend/internal/server/api_contract_test.go @@ -572,6 +572,8 @@ func TestAPIContracts(t *testing.T) { "cache_read_tokens": 2, "cache_creation_5m_tokens": 0, "cache_creation_1h_tokens": 0, + "image_input_tokens": 0, + "image_input_cost": 0, "input_cost": 0, "output_cost": 0, "cache_creation_cost": 0, @@ -927,7 +929,8 @@ func TestAPIContracts(t *testing.T) { "wechat_connect_scopes": "snsapi_login", "withdrawal_management_enabled": true, "withdrawal_rate_limit_window_days": 1, - "withdrawal_rate_limit_max": 0 + "withdrawal_rate_limit_max": 0, + "withdrawal_rate_limit_exempt_amount": 500 } }`, }, @@ -1220,7 +1223,8 @@ func TestAPIContracts(t *testing.T) { "user_private_group_commission_rate": 0.005, "withdrawal_management_enabled": true, "withdrawal_rate_limit_window_days": 1, - "withdrawal_rate_limit_max": 0 + "withdrawal_rate_limit_max": 0, + "withdrawal_rate_limit_exempt_amount": 500 } }`, }, @@ -2493,7 +2497,7 @@ func (r *stubUsageLogRepo) GetAPIKeyDashboardStats(ctx context.Context, apiKeyID return nil, errors.New("not implemented") } -func (r *stubUsageLogRepo) GetUserUsageTrendByUserID(ctx context.Context, userID int64, startTime, endTime time.Time, granularity string) ([]usagestats.TrendDataPoint, error) { +func (r *stubUsageLogRepo) GetUserUsageTrendByUserID(ctx context.Context, userID int64, startTime, endTime time.Time, granularity string, location *time.Location) ([]usagestats.TrendDataPoint, error) { return nil, errors.New("not implemented") } @@ -2501,7 +2505,7 @@ func (r *stubUsageLogRepo) GetUserModelStats(ctx context.Context, userID int64, return nil, errors.New("not implemented") } -func (r *stubUsageLogRepo) GetUserAccountSharingDashboard(ctx context.Context, userID int64, startTime, endTime time.Time, granularity string) (*usagestats.AccountSharingDashboardStats, error) { +func (r *stubUsageLogRepo) GetUserAccountSharingDashboard(ctx context.Context, userID int64, startTime, endTime time.Time, granularity string, location *time.Location) (*usagestats.AccountSharingDashboardStats, error) { return nil, errors.New("not implemented") } diff --git a/backend/internal/server/http.go b/backend/internal/server/http.go index aa7888b73..db4ed3c70 100644 --- a/backend/internal/server/http.go +++ b/backend/internal/server/http.go @@ -6,6 +6,7 @@ import ( "log" "log/slog" "net/http" + "strings" "time" "github.com/Wei-Shaw/sub2api/internal/config" @@ -45,18 +46,7 @@ func ProvideRouter( r := gin.New() r.Use(middleware2.Recovery()) - if len(cfg.Server.TrustedProxies) > 0 { - if err := r.SetTrustedProxies(cfg.Server.TrustedProxies); err != nil { - log.Printf("Failed to set trusted proxies: %v", err) - } - } else { - if err := r.SetTrustedProxies(nil); err != nil { - log.Printf("Failed to disable trusted proxies: %v", err) - } - if cfg.Server.Mode == "release" { - log.Printf("Warning: server.trusted_proxies is empty in release mode; client IP trust chain is disabled") - } - } + configureClientIPResolution(r, cfg.Server, cfg.Security) // Wire up websearch Manager builder so it initializes on startup and rebuilds on config save. settingService.SetWebSearchManagerBuilder(context.Background(), func(cfg *service.WebSearchEmulationConfig, proxyURLs map[int64]string) { @@ -97,6 +87,60 @@ func ProvideRouter( return SetupRouter(r, handlers, jwtAuth, adminAuth, apiKeyAuth, apiKeyService, subscriptionService, opsService, settingService, cfg, redisClient) } +var standardForwardedClientIPHeaders = []string{ + "CF-Connecting-IP", + "X-Forwarded-For", + "X-Real-IP", +} + +func configureClientIPResolution(r *gin.Engine, serverCfg config.ServerConfig, securityCfg config.SecurityConfig) { + customHeaders, err := config.NormalizeForwardedClientIPHeaders(securityCfg.ForwardedClientIPHeaders) + if err != nil { + // Config.Load 会提前拒绝非法配置；这里仍然关闭自定义头，避免直接构造 Config + // 的测试或嵌入调用意外扩大信任边界。 + log.Printf("Ignoring invalid security.forwarded_client_ip_headers: %v", err) + customHeaders = nil + } + r.RemoteIPHeaders = mergeForwardedClientIPHeaders(customHeaders, standardForwardedClientIPHeaders) + + if len(serverCfg.TrustedProxies) == 0 { + if err := r.SetTrustedProxies(nil); err != nil { + log.Printf("Failed to disable trusted proxies: %v", err) + } + if serverCfg.Mode == "release" { + log.Printf("Warning: server.trusted_proxies is empty in release mode; forwarded client IP headers are disabled") + } + return + } + + if err := r.SetTrustedProxies(serverCfg.TrustedProxies); err != nil { + log.Printf("Failed to set trusted proxies, disabling forwarded client IP headers: %v", err) + if disableErr := r.SetTrustedProxies(nil); disableErr != nil { + log.Printf("Failed to disable trusted proxies after invalid configuration: %v", disableErr) + } + } +} + +func mergeForwardedClientIPHeaders(groups ...[]string) []string { + count := 0 + for _, group := range groups { + count += len(group) + } + merged := make([]string, 0, count) + seen := make(map[string]struct{}, count) + for _, group := range groups { + for _, header := range group { + key := strings.ToLower(header) + if _, exists := seen[key]; exists { + continue + } + seen[key] = struct{}{} + merged = append(merged, header) + } + } + return merged +} + // ProvideHTTPServer 提供 HTTP 服务器 func ProvideHTTPServer(cfg *config.Config, router *gin.Engine) *http.Server { httpHandler := http.Handler(router) diff --git a/backend/internal/server/http_client_ip_test.go b/backend/internal/server/http_client_ip_test.go new file mode 100644 index 000000000..ab00c5421 --- /dev/null +++ b/backend/internal/server/http_client_ip_test.go @@ -0,0 +1,100 @@ +//go:build unit + +package server + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/ip" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestConfigureClientIPResolution(t *testing.T) { + gin.SetMode(gin.TestMode) + + tests := []struct { + name string + serverConfig config.ServerConfig + securityConfig config.SecurityConfig + remoteAddr string + headers map[string]string + want string + }{ + { + name: "untrusted peer cannot spoof standard or custom headers", + remoteAddr: "9.9.9.9:12345", + serverConfig: config.ServerConfig{}, + securityConfig: config.SecurityConfig{ + ForwardedClientIPHeaders: []string{"True-Client-IP"}, + }, + headers: map[string]string{ + "True-Client-IP": "1.1.1.1", + "X-Forwarded-For": "2.2.2.2", + "X-Real-IP": "3.3.3.3", + }, + want: "9.9.9.9", + }, + { + name: "trusted proxy accepts explicit custom header", + remoteAddr: "10.0.0.5:12345", + serverConfig: config.ServerConfig{ + TrustedProxies: []string{"10.0.0.0/8"}, + }, + securityConfig: config.SecurityConfig{ + ForwardedClientIPHeaders: []string{"True-Client-IP"}, + }, + headers: map[string]string{ + "True-Client-IP": "1.1.1.1", + "X-Forwarded-For": "2.2.2.2", + }, + want: "1.1.1.1", + }, + { + name: "trusted multi-hop xff skips trusted proxy hops", + remoteAddr: "10.0.0.5:12345", + serverConfig: config.ServerConfig{ + TrustedProxies: []string{"10.0.0.0/8"}, + }, + headers: map[string]string{ + "X-Forwarded-For": "198.51.100.10, 10.0.0.6", + }, + want: "198.51.100.10", + }, + { + name: "invalid trusted proxy configuration fails closed", + remoteAddr: "9.9.9.9:12345", + serverConfig: config.ServerConfig{ + TrustedProxies: []string{"not-a-cidr"}, + }, + headers: map[string]string{ + "X-Forwarded-For": "1.1.1.1", + }, + want: "9.9.9.9", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + router := gin.New() + configureClientIPResolution(router, test.serverConfig, test.securityConfig) + router.GET("/client-ip", func(c *gin.Context) { + c.String(http.StatusOK, ip.GetSecurityClientIP(c)) + }) + + request := httptest.NewRequest(http.MethodGet, "/client-ip", nil) + request.RemoteAddr = test.remoteAddr + for name, value := range test.headers { + request.Header.Set(name, value) + } + response := httptest.NewRecorder() + router.ServeHTTP(response, request) + + require.Equal(t, http.StatusOK, response.Code) + require.Equal(t, test.want, response.Body.String()) + }) + } +} diff --git a/backend/internal/server/middleware/api_key_auth.go b/backend/internal/server/middleware/api_key_auth.go index 852eaa85d..f98d4fd0c 100644 --- a/backend/internal/server/middleware/api_key_auth.go +++ b/backend/internal/server/middleware/api_key_auth.go @@ -3,6 +3,7 @@ package middleware import ( "context" "errors" + "net/http" "strings" "github.com/Wei-Shaw/sub2api/internal/config" @@ -13,6 +14,8 @@ import ( "github.com/gin-gonic/gin" ) +const maxAPIKeyAuthorizationHeaderBytes = service.MaxAPIKeyCredentialBytes + 128 + // NewAPIKeyAuthMiddleware 创建 API Key 认证中间件 func NewAPIKeyAuthMiddleware(apiKeyService *service.APIKeyService, subscriptionService *service.SubscriptionService, cfg *config.Config) APIKeyAuthMiddleware { return APIKeyAuthMiddleware(apiKeyAuthWithSubscription(apiKeyService, subscriptionService, cfg)) @@ -28,6 +31,10 @@ func NewAPIKeyAuthMiddleware(apiKeyService *service.APIKeyService, subscriptionS func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscriptionService *service.SubscriptionService, cfg *config.Config) gin.HandlerFunc { return func(c *gin.Context) { // ── 1. 提取 API Key ────────────────────────────────────────── + if apiKeyHeadersTooLarge(c) { + AbortWithError(c, http.StatusUnauthorized, "INVALID_API_KEY", "Invalid API key") + return + } queryKey := strings.TrimSpace(c.Query("key")) queryApiKey := strings.TrimSpace(c.Query("api_key")) @@ -89,7 +96,7 @@ func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscripti // 检查 IP 限制（白名单/黑名单） // 注意：错误信息故意模糊，避免暴露具体的 IP 限制机制 if len(apiKey.IPWhitelist) > 0 || len(apiKey.IPBlacklist) > 0 { - clientIP := ip.GetTrustedClientIP(c) + clientIP := ip.GetSecurityClientIP(c) allowed, _ := ip.CheckIPRestrictionWithCompiledRules(clientIP, apiKey.CompiledIPWhitelist, apiKey.CompiledIPBlacklist) if !allowed { AbortWithError(c, 403, "ACCESS_DENIED", "Access denied") @@ -156,7 +163,7 @@ func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscripti // Key 状态检查 switch apiKey.Status { case service.StatusAPIKeyQuotaExhausted: - AbortWithError(c, 429, "API_KEY_QUOTA_EXHAUSTED", "API key 额度已用完") + abortWithAPIKeyQuotaError(c) return case service.StatusAPIKeyExpired: AbortWithError(c, 403, "API_KEY_EXPIRED", "API key 已过期") @@ -169,7 +176,7 @@ func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscripti return } if apiKey.IsQuotaExhausted() { - AbortWithError(c, 429, "API_KEY_QUOTA_EXHAUSTED", "API key 额度已用完") + abortWithAPIKeyQuotaError(c) return } @@ -225,6 +232,43 @@ func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscripti } } +func abortWithAPIKeyQuotaError(c *gin.Context) { + const message = "API key 额度已用完" + if isOpenAICompatibleAPIKeyRequest(c) { + abortWithOpenAIQuotaError(c, http.StatusTooManyRequests, message) + return + } + AbortWithError(c, http.StatusTooManyRequests, "API_KEY_QUOTA_EXHAUSTED", message) +} + +func isOpenAICompatibleAPIKeyRequest(c *gin.Context) bool { + if c == nil || c.Request == nil || c.Request.URL == nil { + return false + } + + path := strings.TrimRight(c.Request.URL.Path, "/") + for _, root := range []string{ + "/v1/responses", + "/openai/v1/responses", + "/responses", + "/backend-api/codex/responses", + } { + if path == root || strings.HasPrefix(path, root+"/") { + return true + } + } + return false +} + +func apiKeyHeadersTooLarge(c *gin.Context) bool { + if c == nil { + return false + } + return len(c.GetHeader("Authorization")) > maxAPIKeyAuthorizationHeaderBytes || + len(c.GetHeader("x-api-key")) > service.MaxAPIKeyCredentialBytes || + len(c.GetHeader("x-goog-api-key")) > service.MaxAPIKeyCredentialBytes +} + // 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_test.go b/backend/internal/server/middleware/api_key_auth_test.go index f956fea2b..20b50132b 100644 --- a/backend/internal/server/middleware/api_key_auth_test.go +++ b/backend/internal/server/middleware/api_key_auth_test.go @@ -4,9 +4,11 @@ package middleware import ( "context" + "encoding/json" "errors" "net/http" "net/http/httptest" + "strings" "testing" "time" @@ -503,12 +505,106 @@ func TestAPIKeyAuthTouchesLastUsedInStandardMode(t *testing.T) { require.Equal(t, 1, touchCalls) } +func TestAPIKeyAuthRejectsOversizedCredentialBeforeLookup(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(gin.HandlerFunc(NewAPIKeyAuthMiddleware(nil, nil, &config.Config{}))) + router.GET("/t", func(c *gin.Context) { c.Status(http.StatusOK) }) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/t", nil) + req.Header.Set("x-api-key", strings.Repeat("a", service.MaxAPIKeyCredentialBytes+1)) + router.ServeHTTP(w, req) + + require.Equal(t, http.StatusUnauthorized, w.Code) + var response ErrorResponse + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &response)) + require.Equal(t, "INVALID_API_KEY", response.Code) +} + +func TestAPIKeyAuthOpenAIQuotaErrorFormat(t *testing.T) { + gin.SetMode(gin.TestMode) + + user := &service.User{ID: 11, Role: service.RoleUser, Status: service.StatusActive, Balance: 10} + group := &service.Group{ID: 8, Platform: service.PlatformOpenAI, Status: service.StatusActive} + apiKey := &service.APIKey{ + ID: 105, UserID: user.ID, Key: "openai-quota-exhausted", Status: service.StatusAPIKeyQuotaExhausted, + User: user, Group: group, GroupID: &group.ID, + } + apiKeyRepo := &stubApiKeyRepo{getByKey: func(ctx context.Context, key string) (*service.APIKey, error) { + if key != apiKey.Key { + return nil, service.ErrAPIKeyNotFound + } + clone := *apiKey + userClone := *user + clone.User = &userClone + return &clone, nil + }} + + cfg := &config.Config{RunMode: config.RunModeStandard} + router := newAuthTestRouter(service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg), nil, cfg) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + req.Header.Set("x-api-key", apiKey.Key) + router.ServeHTTP(w, req) + + require.Equal(t, http.StatusTooManyRequests, w.Code) + var response struct { + Error struct { + Message string `json:"message"` + Type string `json:"type"` + Param *string `json:"param"` + Code string `json:"code"` + } `json:"error"` + } + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &response)) + require.Equal(t, "API key 额度已用完", response.Error.Message) + require.Equal(t, "insufficient_quota", response.Error.Type) + require.Nil(t, response.Error.Param) + require.Equal(t, "insufficient_quota", response.Error.Code) +} + +func TestAPIKeyAuthQuotaErrorKeepsLegacyFormatOutsideResponses(t *testing.T) { + gin.SetMode(gin.TestMode) + + user := &service.User{ID: 11, Role: service.RoleUser, Status: service.StatusActive, Balance: 10} + group := &service.Group{ID: 8, Platform: service.PlatformOpenAI, Status: service.StatusActive} + apiKey := &service.APIKey{ + ID: 106, UserID: user.ID, Key: "legacy-quota-exhausted", Status: service.StatusAPIKeyQuotaExhausted, + User: user, Group: group, GroupID: &group.ID, + } + apiKeyRepo := &stubApiKeyRepo{getByKey: func(ctx context.Context, key string) (*service.APIKey, error) { + if key != apiKey.Key { + return nil, service.ErrAPIKeyNotFound + } + clone := *apiKey + userClone := *user + clone.User = &userClone + return &clone, nil + }} + + cfg := &config.Config{RunMode: config.RunModeStandard} + router := newAuthTestRouter(service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg), nil, cfg) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/v1/messages", nil) + req.Header.Set("x-api-key", apiKey.Key) + router.ServeHTTP(w, req) + + require.Equal(t, http.StatusTooManyRequests, w.Code) + var response ErrorResponse + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &response)) + require.Equal(t, "API_KEY_QUOTA_EXHAUSTED", response.Code) + require.Equal(t, "API key 额度已用完", response.Message) +} + 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) { c.JSON(http.StatusOK, gin.H{"ok": true}) }) + router.POST("/v1/responses", func(c *gin.Context) { c.Status(http.StatusOK) }) + router.POST("/v1/messages", func(c *gin.Context) { c.Status(http.StatusOK) }) return router } diff --git a/backend/internal/server/middleware/middleware.go b/backend/internal/server/middleware/middleware.go index 27985cf8b..1dcd6141b 100644 --- a/backend/internal/server/middleware/middleware.go +++ b/backend/internal/server/middleware/middleware.go @@ -75,6 +75,19 @@ func AbortWithError(c *gin.Context, statusCode int, code, message string) { c.Abort() } +// abortWithOpenAIQuotaError writes the OpenAI-compatible insufficient quota response. +func abortWithOpenAIQuotaError(c *gin.Context, statusCode int, message string) { + c.JSON(statusCode, gin.H{ + "error": gin.H{ + "message": message, + "type": "insufficient_quota", + "param": nil, + "code": "insufficient_quota", + }, + }) + c.Abort() +} + // ────────────────────────────────────────────────────────── // RequireGroupAssignment — 未分组 Key 拦截中间件 // ────────────────────────────────────────────────────────── diff --git a/backend/internal/server/routes/gateway.go b/backend/internal/server/routes/gateway.go index 5e006e1e0..1745f2add 100644 --- a/backend/internal/server/routes/gateway.go +++ b/backend/internal/server/routes/gateway.go @@ -101,6 +101,18 @@ func RegisterGatewayRoutes( }, }) } + videoContentHandler := func(c *gin.Context) { + if getGroupPlatform(c) == service.PlatformGrok { + h.OpenAIGateway.GrokVideoContent(c) + return + } + c.JSON(http.StatusNotFound, gin.H{ + "error": gin.H{ + "type": "not_found_error", + "message": "Videos API is not supported for this platform", + }, + }) + } // API网关（Claude API兼容） gateway := r.Group("/v1") @@ -168,6 +180,7 @@ func RegisterGatewayRoutes( gateway.POST("/videos/edits", videoEditHandler) gateway.POST("/videos/extensions", videoExtensionHandler) gateway.GET("/videos/:request_id", videoStatusHandler) + gateway.GET("/videos/:request_id/content", videoContentHandler) } // Gemini 原生 API 兼容层（Gemini SDK/CLI 直连） @@ -221,6 +234,7 @@ func RegisterGatewayRoutes( r.POST("/videos/edits", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, videoEditHandler) r.POST("/videos/extensions", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, videoExtensionHandler) r.GET("/videos/:request_id", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, videoStatusHandler) + r.GET("/videos/:request_id/content", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, videoContentHandler) // Antigravity 模型列表 r.GET("/antigravity/models", gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, h.Gateway.AntigravityModels) diff --git a/backend/internal/server/routes/gateway_test.go b/backend/internal/server/routes/gateway_test.go index d7f0b86d0..a53299922 100644 --- a/backend/internal/server/routes/gateway_test.go +++ b/backend/internal/server/routes/gateway_test.go @@ -122,6 +122,22 @@ func TestGatewayRoutesGrokVideoMutationPathsAreRegistered(t *testing.T) { } } +func TestGatewayRoutesGrokVideoContentPathsAreRegistered(t *testing.T) { + router := newGatewayRoutesTestRouter(service.PlatformGrok) + registered := make(map[string]bool) + for _, route := range router.Routes() { + if route.Method == http.MethodGet { + registered[route.Path] = true + } + } + for _, path := range []string{ + "/v1/videos/:request_id/content", + "/videos/:request_id/content", + } { + require.True(t, registered[path], "GET %s should be registered", path) + } +} + func TestGatewayRoutesAlphaSearchRejectsNonOpenAIGroup(t *testing.T) { router := newGatewayRoutesTestRouter(service.PlatformGrok) req := httptest.NewRequest(http.MethodPost, "/v1/alpha/search", strings.NewReader(`{"model":"gpt-5.6-sol"}`)) diff --git a/backend/internal/service/account.go b/backend/internal/service/account.go index c8a717d12..d3c92e88d 100644 --- a/backend/internal/service/account.go +++ b/backend/internal/service/account.go @@ -8,6 +8,7 @@ import ( "fmt" "hash/fnv" "log/slog" + "net/http" "reflect" "sort" "strconv" @@ -89,6 +90,23 @@ type Account struct { headerOverrideCacheRawSig uint64 } +// OpenAIEndpointCapability identifies an endpoint-specific requirement that +// must survive scheduler cache hydration and the final pre-dispatch recheck. +type OpenAIEndpointCapability string + +const ( + // OpenAIEndpointCapabilityGrokMediaGeneration keeps new image/video + // generation requests away from Grok OAuth accounts without positive paid + // entitlement evidence. Video status/content lookups intentionally do not + // require this capability so existing tasks remain queryable. + OpenAIEndpointCapabilityGrokMediaGeneration OpenAIEndpointCapability = "grok_media_generation" +) + +// GrokMediaEligibleExtraKey is an optional operator override in accounts.extra. +// A boolean true/false takes precedence over provider observations; absent, +// null, or malformed values do not override observed billing state. +const GrokMediaEligibleExtraKey = "grok_media_eligible" + const ( AccountShareModePrivate = "private" AccountShareModePublic = "public" @@ -1653,13 +1671,25 @@ func (a *Account) IsOpenAIOAuth() bool { return a.IsOpenAI() && a.Type == AccountTypeOAuth } +// IsOpenAIAgentIdentityCredentials reports whether credentials select the +// Codex Agent Identity authentication mode. Platform and account-type checks +// remain the caller's responsibility so this helper can also be used while a +// create/import request is still being validated. +func IsOpenAIAgentIdentityCredentials(credentials map[string]any) bool { + if len(credentials) == 0 { + return false + } + authMode, ok := credentials["auth_mode"].(string) + return ok && strings.EqualFold(strings.TrimSpace(authMode), OpenAIAuthModeAgentIdentity) +} + // IsOpenAIAgentIdentity reports whether the account uses Codex Agent Identity // credentials instead of a refreshable OpenAI OAuth token. func (a *Account) IsOpenAIAgentIdentity() bool { if a == nil || !a.IsOpenAIOAuth() { return false } - return strings.EqualFold(strings.TrimSpace(a.GetCredential("auth_mode")), OpenAIAuthModeAgentIdentity) + return IsOpenAIAgentIdentityCredentials(a.Credentials) } // ValidateOpenAIAgentIdentityPrivateKey validates the base64-encoded PKCS#8 @@ -2121,6 +2151,79 @@ func (a *Account) GetOpenAISessionID() string { return strings.TrimSpace(a.GetExtraString("openai_session_id")) } +// SupportsOpenAIEndpointCapability reports whether an account can remain a +// scheduler candidate for an endpoint-specific request. Unobserved Grok OAuth +// accounts remain candidates only so the request path can run the billing +// probe; forwarding still fails closed unless that probe yields paid evidence. +func (a *Account) SupportsOpenAIEndpointCapability(capability OpenAIEndpointCapability) bool { + if a == nil { + return false + } + if capability == "" { + return true + } + if !a.IsOpenAICompatible() { + return false + } + switch capability { + case OpenAIEndpointCapabilityGrokMediaGeneration: + if !a.IsGrok() { + return false + } + eligible, reason := a.GrokMediaGenerationEligibility() + return eligible || reason == "billing_unobserved" + default: + return false + } +} + +// GrokMediaGenerationEligibility reports whether a Grok account may receive +// a new image/video generation request. OAuth accounts fail closed unless +// billing observations provide positive paid-entitlement evidence. +func (a *Account) GrokMediaGenerationEligibility() (bool, string) { + if a == nil || !a.IsGrok() { + return false, "not_grok" + } + if override, ok := grokMediaEligibilityOverride(a.Extra); ok { + if override { + return true, "override_enabled" + } + return false, "override_disabled" + } + if a.Type != AccountTypeOAuth { + return true, "non_oauth" + } + + billing, err := grokBillingSnapshotFromExtra(a.Extra) + if err != nil || billing == nil { + return false, "billing_unobserved" + } + if billing.StatusCode == http.StatusForbidden || + billing.WeeklyStatusCode == http.StatusForbidden || + billing.MonthlyStatusCode == http.StatusForbidden { + return false, "billing_forbidden" + } + if isKnownGrokFreeAccount(a) { + return false, "billing_free_tier" + } + if !grokBillingHasAuthoritativeQuota(billing) { + return false, "billing_inconclusive" + } + return true, "eligible" +} + +func grokMediaEligibilityOverride(extra map[string]any) (bool, bool) { + if extra == nil { + return false, false + } + raw, exists := extra[GrokMediaEligibleExtraKey] + if !exists || raw == nil { + return false, false + } + value, ok := raw.(bool) + return value, ok +} + func (a *Account) SupportsOpenAIImageCapability(capability OpenAIImagesCapability) bool { if !a.IsOpenAI() { return false diff --git a/backend/internal/service/account_credential_import.go b/backend/internal/service/account_credential_import.go index 80359d109..59007e9f4 100644 --- a/backend/internal/service/account_credential_import.go +++ b/backend/internal/service/account_credential_import.go @@ -60,6 +60,7 @@ type AccountCredentialImportError struct { type AccountCredentialImportResult struct { Total int `json:"total"` Created int `json:"created"` + Updated int `json:"updated"` Failed int `json:"failed"` Errors []AccountCredentialImportError `json:"errors"` } @@ -244,6 +245,9 @@ func accountCredentialImportSourcesFromValue(value any) ([]AccountCredentialImpo } func accountCredentialImportSourceFromMap(item map[string]any) (AccountCredentialImportSource, error) { + if source, handled, err := accountCredentialImportSourceFromAgentIdentityAccountEnvelope(item); handled || err != nil { + return source, err + } if source, handled, err := accountCredentialImportSourceFromAgentIdentity(item); handled || err != nil { return source, err } @@ -375,6 +379,76 @@ func accountCredentialImportSourceFromMap(item map[string]any) (AccountCredentia return AccountCredentialImportSource{}, fmt.Errorf("unsupported credential import item") } +func accountCredentialImportSourceFromAgentIdentityAccountEnvelope(item map[string]any) (AccountCredentialImportSource, bool, error) { + credentialsValue, hasCredentials := importAnyField(item, "credentials") + if !hasCredentials { + return AccountCredentialImportSource{}, false, nil + } + credentials, ok := credentialsValue.(map[string]any) + if !ok { + return AccountCredentialImportSource{}, false, nil + } + + authMode := importStringField(credentials, "auth_mode", "authMode") + _, hasIdentity := importAnyField(credentials, "agent_identity", "agentIdentity") + if !hasIdentity && !strings.EqualFold(strings.TrimSpace(authMode), OpenAIAuthModeAgentIdentity) { + return AccountCredentialImportSource{}, false, nil + } + + if outerAuthMode := importStringField(item, "auth_mode", "authMode"); outerAuthMode != "" { + return AccountCredentialImportSource{}, true, fmt.Errorf("Agent Identity auth_mode must be declared only inside credentials") + } + if _, hasOuterIdentity := importAnyField(item, "agent_identity", "agentIdentity"); hasOuterIdentity { + return AccountCredentialImportSource{}, true, fmt.Errorf("Agent Identity must not be declared in both the account and credentials") + } + declaredPlatform := importStringField(item, "platform", "provider", "service") + if declaredPlatform != "" && normalizeCredentialImportPlatform(declaredPlatform) != PlatformOpenAI { + return AccountCredentialImportSource{}, true, fmt.Errorf("Agent Identity account platform must be OpenAI") + } + accountType := strings.ToLower(strings.TrimSpace(importStringField(item, "type", "account_type", "accountType"))) + if accountType != "" && accountType != AccountTypeOAuth { + return AccountCredentialImportSource{}, true, fmt.Errorf("Agent Identity account type must be OAuth") + } + + outerSafety := copyImportMap(item) + removeImportMapField(outerSafety, "credentials") + if field, found := findOAuthTokenCredentialImportField(outerSafety); found { + return AccountCredentialImportSource{}, true, fmt.Errorf("Agent Identity must not include OAuth token field: %s", field) + } + if field, found := findDisallowedCredentialImportField(outerSafety); found { + return AccountCredentialImportSource{}, true, fmt.Errorf("disallowed credential field: %s", field) + } + + // Legacy account exports can contain an OAuth id_token alongside Agent + // Identity credentials. Agent Identity authentication does not use it, so + // discard only this direct envelope field before the strict recursive token + // scan. Access/refresh tokens and every nested token field remain rejected. + identityInput := copyImportMap(credentials) + if idTokenValue, hasIDToken := importAnyField(identityInput, "id_token", "idToken"); hasIDToken { + if _, valid := idTokenValue.(string); !valid { + return AccountCredentialImportSource{}, true, fmt.Errorf("Agent Identity account id_token must be a string") + } + removeImportMapField(identityInput, "id_token") + removeImportMapField(identityInput, "idToken") + } + + source, handled, err := accountCredentialImportSourceFromAgentIdentity(identityInput) + if err != nil { + return AccountCredentialImportSource{}, true, err + } + if !handled { + return AccountCredentialImportSource{}, true, fmt.Errorf("Agent Identity credentials are invalid") + } + + source.Name = credentialImportFirstNonEmptyString( + importStringField(item, "name", "label"), + source.Name, + ) + source.Notes = importOptionalStringField(item, "notes", "note", "description") + source.Extra = importMapField(item, "extra", "metadata") + return source, true, nil +} + func accountCredentialImportSourceFromAgentIdentity(item map[string]any) (AccountCredentialImportSource, bool, error) { authMode := importStringField(item, "auth_mode", "authMode") identityValue, hasIdentity := importAnyField(item, "agent_identity", "agentIdentity") @@ -385,6 +459,9 @@ func accountCredentialImportSourceFromAgentIdentity(item map[string]any) (Accoun if authMode != "" && !isAgentAuthMode { return AccountCredentialImportSource{}, true, fmt.Errorf("Agent Identity auth_mode is invalid") } + if field, found := findOAuthTokenCredentialImportField(item); found { + return AccountCredentialImportSource{}, true, fmt.Errorf("Agent Identity must not include OAuth token field: %s", field) + } identity := item if hasIdentity { @@ -585,6 +662,10 @@ func findDisallowedCredentialImportField(value any) (string, bool) { }) } +func findOAuthTokenCredentialImportField(value any) (string, bool) { + return findOAuthTokenCredentialContent(value) +} + func importMapField(values map[string]any, keys ...string) map[string]any { value, ok := importAnyField(values, keys...) if !ok { diff --git a/backend/internal/service/account_credential_import_test.go b/backend/internal/service/account_credential_import_test.go index c7fe3ceb7..a6129a9bd 100644 --- a/backend/internal/service/account_credential_import_test.go +++ b/backend/internal/service/account_credential_import_test.go @@ -67,6 +67,114 @@ func TestAccountCredentialImportSupportsAgentIdentitySchemas(t *testing.T) { } } +func TestAccountCredentialImportSupportsAgentIdentityAccountExportEnvelope(t *testing.T) { + privateKey := testAgentIdentityPrivateKey(t) + content := `{ + "type":"sub2api-data", + "version":1, + "exported_at":"2026-07-22T00:00:00Z", + "proxies":[], + "accounts":[{ + "name":"exported agent", + "platform":"openai", + "type":"oauth", + "credentials":{ + "account_id":"legacy-account", + "agent_private_key":"` + privateKey + `", + "agent_runtime_id":"runtime-export", + "auth_mode":"agentIdentity", + "chatgpt_account_id":"team-export", + "chatgpt_user_id":"user-export", + "id_token":"legacy-id-token-must-not-be-stored", + "workspace_id":"workspace-export" + }, + "extra":{"email":"agent@example.com","email_key":"mail-key"}, + "concurrency":1, + "priority":50 + }] + }` + + sources, errs := ParseAccountCredentialImportContents([]string{content}) + if len(errs) != 0 { + t.Fatalf("errs = %#v, want none", errs) + } + if len(sources) != 1 { + t.Fatalf("sources len = %d, want 1", len(sources)) + } + source := sources[0] + if source.Kind != AccountCredentialImportKindOpenAIAgentIdentity { + t.Fatalf("kind = %s, want %s", source.Kind, AccountCredentialImportKindOpenAIAgentIdentity) + } + if source.Name != "exported agent" { + t.Fatalf("name = %q, want exported agent", source.Name) + } + if source.Credentials["agent_runtime_id"] != "runtime-export" || source.Credentials["chatgpt_account_id"] != "team-export" { + t.Fatalf("credentials identifiers were not preserved: %#v", source.Credentials) + } + for _, field := range []string{"access_token", "refresh_token", "id_token", "account_id", "workspace_id"} { + if _, ok := source.Credentials[field]; ok { + t.Fatalf("credential field %q must not be retained", field) + } + } + if source.Extra["email_key"] != "mail-key" { + t.Fatalf("safe extra metadata was not preserved: %#v", source.Extra) + } + if err := validateOwnedAccountSourceForPlatform(source.Platform, AccountTypeOAuth, source.Credentials, source.Extra); err != nil { + t.Fatalf("parsed envelope failed owned-account validation: %v", err) + } +} + +func TestAccountCredentialImportAgentIdentityAccountExportEnvelopeRejectsUnsafeTokens(t *testing.T) { + privateKey := testAgentIdentityPrivateKey(t) + tests := []struct { + name string + injected string + want string + }{ + {name: "access token", injected: `"access_token":"must-reject"`, want: "access_token"}, + {name: "refresh token", injected: `"refreshToken":"must-reject"`, want: "refreshToken"}, + {name: "nested id token", injected: `"metadata":{"id_token":"must-reject"}`, want: "id_token"}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + content := `{"accounts":[{"name":"agent","platform":"openai","type":"oauth","credentials":{` + + `"auth_mode":"agentIdentity","agent_runtime_id":"runtime-1","agent_private_key":"` + privateKey + `",` + + `"chatgpt_account_id":"team-1",` + test.injected + `}}]}` + _, errs := ParseAccountCredentialImportContents([]string{content}) + if len(errs) != 1 || !strings.Contains(errs[0].Message, test.want) { + t.Fatalf("errs = %#v, want one error containing %q", errs, test.want) + } + }) + } +} + +func TestAccountCredentialImportAgentIdentityAccountExportEnvelopeRejectsInvalidMetadata(t *testing.T) { + privateKey := testAgentIdentityPrivateKey(t) + tests := []struct { + name string + outerFields string + want string + }{ + {name: "wrong platform", outerFields: `"platform":"anthropic","type":"oauth",`, want: "platform must be OpenAI"}, + {name: "wrong account type", outerFields: `"platform":"openai","type":"api_key",`, want: "type must be OAuth"}, + {name: "duplicate auth mode", outerFields: `"platform":"openai","type":"oauth","auth_mode":"agentIdentity",`, want: "auth_mode must be declared only inside credentials"}, + {name: "duplicate identity", outerFields: `"platform":"openai","type":"oauth","agent_identity":{},`, want: "must not be declared in both"}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + content := `{"accounts":[{` + test.outerFields + `"credentials":{` + + `"auth_mode":"agentIdentity","agent_runtime_id":"runtime-1","agent_private_key":"` + privateKey + `",` + + `"chatgpt_account_id":"team-1"}}]}` + _, errs := ParseAccountCredentialImportContents([]string{content}) + if len(errs) != 1 || !strings.Contains(errs[0].Message, test.want) { + t.Fatalf("errs = %#v, want one error containing %q", errs, test.want) + } + }) + } +} + func TestAccountCredentialImportRejectsInvalidAgentIdentity(t *testing.T) { tests := []struct { name string @@ -88,6 +196,26 @@ func TestAccountCredentialImportRejectsInvalidAgentIdentity(t *testing.T) { content: `{"agent_identity":{"agent_runtime_id":"runtime-1","agent_private_key":"` + testAgentIdentityPrivateKey(t) + `","base_url":"https://evil.example"}}`, want: "disallowed credential field: base_url", }, + { + name: "top-level OAuth token", + content: `{"auth_mode":"agentIdentity","agent_runtime_id":"runtime-1","agent_private_key":"` + testAgentIdentityPrivateKey(t) + `","access_token":"must-reject"}`, + want: "must not include OAuth token field: access_token", + }, + { + name: "top-level ID token", + content: `{"auth_mode":"agentIdentity","agent_runtime_id":"runtime-1","agent_private_key":"` + testAgentIdentityPrivateKey(t) + `","id_token":"must-reject"}`, + want: "must not include OAuth token field: id_token", + }, + { + name: "nested OAuth token", + content: `{"agent_identity":{"agent_runtime_id":"runtime-1","agent_private_key":"` + testAgentIdentityPrivateKey(t) + `","tokens":[{"refreshToken":"must-reject"}]}}`, + want: "must not include OAuth token field: refreshToken", + }, + { + name: "nested OAuth token normalized key", + content: `{"agent_identity":{"agent_runtime_id":"runtime-1","agent_private_key":"` + testAgentIdentityPrivateKey(t) + `","metadata":{" ID_TOKEN ":"must-reject"}}}`, + want: "must not include OAuth token field: ID_TOKEN ", + }, } for _, test := range tests { diff --git a/backend/internal/service/account_credential_safety.go b/backend/internal/service/account_credential_safety.go index 6baad2cae..508716708 100644 --- a/backend/internal/service/account_credential_safety.go +++ b/backend/internal/service/account_credential_safety.go @@ -6,6 +6,8 @@ type credentialSafetyOptions struct { AllowClaudeSessionKeyFields bool AllowOAuthTokenValues bool AllowOAuthMetadataURLs bool + DisallowOAuthTokenFields bool + OnlyOAuthTokenFields bool } func findDisallowedCredentialContent(value any, opts credentialSafetyOptions) (string, bool) { @@ -47,6 +49,12 @@ func findDisallowedCredentialContentAt(value any, parentKey string, opts credent } func isDisallowedCredentialSafetyFieldKey(normalizedKey string, opts credentialSafetyOptions) bool { + if isOAuthTokenCredentialSafetyFieldKey(normalizedKey) { + return opts.DisallowOAuthTokenFields || opts.OnlyOAuthTokenFields + } + if opts.OnlyOAuthTokenFields { + return false + } switch normalizedKey { case "api_key", "apikey", @@ -99,7 +107,19 @@ func isDisallowedCredentialSafetyFieldKey(normalizedKey string, opts credentialS return false } +func isOAuthTokenCredentialSafetyFieldKey(normalizedKey string) bool { + switch normalizedKey { + case "access_token", "accesstoken", "refresh_token", "refreshtoken", "id_token", "idtoken": + return true + default: + return false + } +} + func disallowedCredentialStringReason(key, value string, opts credentialSafetyOptions) (string, bool) { + if opts.OnlyOAuthTokenFields { + return "", false + } text := strings.TrimSpace(value) if text == "" { return "", false @@ -124,6 +144,10 @@ func disallowedCredentialStringReason(key, value string, opts credentialSafetyOp return "", false } +func findOAuthTokenCredentialContent(value any) (string, bool) { + return findDisallowedCredentialContent(value, credentialSafetyOptions{OnlyOAuthTokenFields: true}) +} + func isAllowedOAuthMetadataURLField(key string) bool { switch normalizeCredentialSafetyKey(key) { case "scope", diff --git a/backend/internal/service/account_grok_managed_extra.go b/backend/internal/service/account_grok_managed_extra.go new file mode 100644 index 000000000..4bbf45667 --- /dev/null +++ b/backend/internal/service/account_grok_managed_extra.go @@ -0,0 +1,85 @@ +package service + +import ( + "encoding/json" + + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" +) + +var ( + ErrOwnedAccountGrokManagedExtraNotAllowed = infraerrors.BadRequest( + "OWNED_ACCOUNT_GROK_MANAGED_EXTRA_NOT_ALLOWED", + "user accounts cannot modify server-managed Grok media eligibility metadata", + ) + ErrGrokBillingSnapshotManaged = infraerrors.BadRequest( + "GROK_BILLING_SNAPSHOT_MANAGED", + "grok_billing_snapshot is server-managed and cannot be set manually", + ) + ErrGrokMediaEligibilityOverrideInvalid = infraerrors.BadRequest( + "GROK_MEDIA_ELIGIBILITY_OVERRIDE_INVALID", + "grok_media_eligible must be a boolean or null", + ) +) + +var ownedAccountGrokManagedExtraKeys = [...]string{ + GrokMediaEligibleExtraKey, + grokBillingExtraKey, +} + +// rejectOwnedAccountGrokManagedExtra prevents user-owned account APIs from +// forging operator overrides or provider observations used by the scheduler. +func rejectOwnedAccountGrokManagedExtra(extra map[string]any) error { + for _, key := range ownedAccountGrokManagedExtraKeys { + if _, exists := extra[key]; exists { + return ErrOwnedAccountGrokManagedExtraNotAllowed.WithMetadata(map[string]string{"field": key}) + } + } + return nil +} + +// preserveOwnedAccountGrokManagedExtra keeps server-managed values on a +// replacement-style user update. Echoing the unchanged value is accepted so +// clients may safely round-trip an account response, while mutations and +// explicit deletion attempts fail closed. +func preserveOwnedAccountGrokManagedExtra(current, next map[string]any) error { + for _, key := range ownedAccountGrokManagedExtraKeys { + requested, provided := next[key] + stored, exists := current[key] + if provided && (!exists || !sameAccountJSONValue(stored, requested)) { + return ErrOwnedAccountGrokManagedExtraNotAllowed.WithMetadata(map[string]string{"field": key}) + } + preserveMapKey(current, next, key) + } + return nil +} + +// validateAdminGrokManagedExtra allows administrators to set only the +// documented boolean eligibility override. Billing observations are written +// exclusively by the quota probe and cannot enter through generic CRUD APIs. +func validateAdminGrokManagedExtra(extra map[string]any) (mediaOverrideProvided bool, err error) { + if _, exists := extra[grokBillingExtraKey]; exists { + return false, ErrGrokBillingSnapshotManaged + } + raw, exists := extra[GrokMediaEligibleExtraKey] + if !exists { + return false, nil + } + if raw != nil { + if _, ok := raw.(bool); !ok { + return false, ErrGrokMediaEligibilityOverrideInvalid + } + } + return true, nil +} + +func sameAccountJSONValue(left, right any) bool { + leftJSON, leftErr := json.Marshal(left) + if leftErr != nil { + return false + } + rightJSON, rightErr := json.Marshal(right) + if rightErr != nil { + return false + } + return string(leftJSON) == string(rightJSON) +} diff --git a/backend/internal/service/account_grok_managed_extra_test.go b/backend/internal/service/account_grok_managed_extra_test.go new file mode 100644 index 000000000..ec14e3113 --- /dev/null +++ b/backend/internal/service/account_grok_managed_extra_test.go @@ -0,0 +1,226 @@ +//go:build unit + +package service + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestNoAvailableOpenAISelectionErrorWrapsSentinel(t *testing.T) { + t.Parallel() + + for _, requestedModel := range []string{"", "grok-4"} { + err := noAvailableOpenAISelectionError(requestedModel, false) + require.True(t, errors.Is(err, ErrNoAvailableAccounts)) + require.Contains(t, err.Error(), "no available OpenAI accounts") + } + require.ErrorIs(t, noAvailableOpenAISelectionError("grok-4", true), ErrNoAvailableCompactAccounts) +} + +func TestAccountServiceOwnedCreateRejectsGrokManagedExtra(t *testing.T) { + t.Parallel() + + for _, key := range []string{GrokMediaEligibleExtraKey, grokBillingExtraKey} { + key := key + t.Run(key, func(t *testing.T) { + t.Parallel() + svc := &AccountService{} + _, err := svc.CreateOwned(context.Background(), 101, CreateAccountRequest{ + Extra: map[string]any{key: true}, + }) + require.ErrorIs(t, err, ErrOwnedAccountGrokManagedExtraNotAllowed) + }) + } +} + +func TestAccountServiceOwnedUpdatePreservesUnchangedGrokManagedExtra(t *testing.T) { + t.Parallel() + + repo := newOwnedAgentIdentityRepoStub() + svc, _ := newOwnedAgentIdentityService(repo) + created, err := svc.ImportOwnedWithResult(context.Background(), 101, ownedAgentIdentityImportRequest(t, "team-managed", "member-managed", "runtime-managed", "team")) + require.NoError(t, err) + + snapshot := map[string]any{"plan": "SuperGrok", "status_code": float64(200)} + repo.accounts[created.Account.ID].Extra = map[string]any{ + GrokMediaEligibleExtraKey: true, + grokBillingExtraKey: snapshot, + "old_config": "remove-me", + } + + t.Run("omitted values are preserved", func(t *testing.T) { + extra := map[string]any{"custom": "next"} + updated, updateErr := svc.UpdateOwned(context.Background(), 101, created.Account.ID, UpdateAccountRequest{Extra: &extra}) + require.NoError(t, updateErr) + require.Equal(t, true, updated.Extra[GrokMediaEligibleExtraKey]) + require.Equal(t, snapshot, updated.Extra[grokBillingExtraKey]) + require.Equal(t, "next", updated.Extra["custom"]) + }) + + t.Run("unchanged echoed values are accepted", func(t *testing.T) { + extra := map[string]any{ + GrokMediaEligibleExtraKey: true, + grokBillingExtraKey: map[string]any{"plan": "SuperGrok", "status_code": float64(200)}, + "custom": "echo", + } + updated, updateErr := svc.UpdateOwned(context.Background(), 101, created.Account.ID, UpdateAccountRequest{Extra: &extra}) + require.NoError(t, updateErr) + require.Equal(t, true, updated.Extra[GrokMediaEligibleExtraKey]) + require.Equal(t, snapshot, updated.Extra[grokBillingExtraKey]) + }) +} + +func TestAccountServiceOwnedUpdateRejectsGrokManagedExtraMutation(t *testing.T) { + t.Parallel() + + for _, testCase := range []struct { + name string + extra map[string]any + }{ + {name: "override mutation", extra: map[string]any{GrokMediaEligibleExtraKey: false}}, + {name: "snapshot deletion", extra: map[string]any{grokBillingExtraKey: nil}}, + } { + testCase := testCase + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + repo := newOwnedAgentIdentityRepoStub() + svc, _ := newOwnedAgentIdentityService(repo) + created, err := svc.ImportOwnedWithResult(context.Background(), 101, ownedAgentIdentityImportRequest(t, "team-"+testCase.name, "member", "runtime", "team")) + require.NoError(t, err) + repo.accounts[created.Account.ID].Extra = map[string]any{ + GrokMediaEligibleExtraKey: true, + grokBillingExtraKey: map[string]any{"status_code": float64(200)}, + } + updatesBefore := repo.updateCount + + _, err = svc.UpdateOwned(context.Background(), 101, created.Account.ID, UpdateAccountRequest{Extra: &testCase.extra}) + + require.ErrorIs(t, err, ErrOwnedAccountGrokManagedExtraNotAllowed) + require.Equal(t, updatesBefore, repo.updateCount) + }) + } +} + +func TestAccountServiceOwnedBulkRejectsGrokManagedExtra(t *testing.T) { + t.Parallel() + + for _, key := range []string{GrokMediaEligibleExtraKey, grokBillingExtraKey} { + key := key + t.Run(key, func(t *testing.T) { + t.Parallel() + svc := &AccountService{} + _, err := svc.BulkUpdateOwned(context.Background(), 101, &BulkUpdateOwnedAccountsInput{ + AccountIDs: []int64{1}, + Extra: map[string]any{key: nil}, + }) + require.ErrorIs(t, err, ErrOwnedAccountGrokManagedExtraNotAllowed) + }) + } +} + +func TestAdminAccountCreateValidatesGrokManagedExtra(t *testing.T) { + t.Parallel() + + svc := &adminServiceImpl{} + for _, testCase := range []struct { + name string + extra map[string]any + wantErr error + }{ + {name: "boolean override", extra: map[string]any{GrokMediaEligibleExtraKey: true}}, + {name: "null override", extra: map[string]any{GrokMediaEligibleExtraKey: nil}}, + {name: "invalid override", extra: map[string]any{GrokMediaEligibleExtraKey: "true"}, wantErr: ErrGrokMediaEligibilityOverrideInvalid}, + {name: "billing snapshot", extra: map[string]any{grokBillingExtraKey: map[string]any{}}, wantErr: ErrGrokBillingSnapshotManaged}, + } { + testCase := testCase + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + account, _, err := svc.prepareAccountCreate(context.Background(), &CreateAccountInput{ + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Extra: testCase.extra, + SkipDefaultGroupBind: true, + }) + if testCase.wantErr != nil { + require.ErrorIs(t, err, testCase.wantErr) + require.Nil(t, account) + return + } + require.NoError(t, err) + require.Contains(t, account.Extra, GrokMediaEligibleExtraKey) + }) + } +} + +func TestAdminAccountUpdatePreservesGrokBillingSnapshot(t *testing.T) { + t.Parallel() + + for _, testCase := range []struct { + name string + extra map[string]any + wantMediaOverride any + }{ + {name: "ordinary update preserves override", extra: map[string]any{"custom": "next"}, wantMediaOverride: true}, + {name: "explicit null clears override", extra: map[string]any{GrokMediaEligibleExtraKey: nil, "custom": "next"}, wantMediaOverride: nil}, + } { + testCase := testCase + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + repo := newOwnedAgentIdentityRepoStub() + repo.accounts[1] = &Account{ + ID: 1, + Platform: PlatformGrok, + Type: AccountTypeOAuth, + AccountLevel: AccountLevelUnknown, + Concurrency: 1, + Status: StatusActive, + Schedulable: true, + Extra: map[string]any{ + GrokMediaEligibleExtraKey: true, + grokBillingExtraKey: map[string]any{"status_code": float64(200)}, + "old_config": "replace-me", + }, + } + svc := &adminServiceImpl{accountRepo: repo} + + updated, err := svc.UpdateAccount(context.Background(), 1, &UpdateAccountInput{Extra: testCase.extra}) + + require.NoError(t, err) + require.Equal(t, testCase.wantMediaOverride, updated.Extra[GrokMediaEligibleExtraKey]) + require.Equal(t, map[string]any{"status_code": float64(200)}, updated.Extra[grokBillingExtraKey]) + require.Equal(t, "next", updated.Extra["custom"]) + }) + } +} + +func TestAdminAccountUpdateAndBulkRejectGrokBillingSnapshotInjection(t *testing.T) { + t.Parallel() + + repo := newOwnedAgentIdentityRepoStub() + repo.accounts[1] = &Account{ID: 1, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusActive} + svc := &adminServiceImpl{accountRepo: repo} + + _, err := svc.UpdateAccount(context.Background(), 1, &UpdateAccountInput{ + Extra: map[string]any{grokBillingExtraKey: map[string]any{"status_code": float64(200)}}, + }) + require.ErrorIs(t, err, ErrGrokBillingSnapshotManaged) + require.Zero(t, repo.updateCount) + + _, err = svc.BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{ + AccountIDs: []int64{1}, + Extra: map[string]any{grokBillingExtraKey: nil}, + }) + require.ErrorIs(t, err, ErrGrokBillingSnapshotManaged) + require.Zero(t, repo.bulkUpdateCalls) + + _, err = svc.BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{ + AccountIDs: []int64{1}, + Extra: map[string]any{GrokMediaEligibleExtraKey: 1}, + }) + require.ErrorIs(t, err, ErrGrokMediaEligibilityOverrideInvalid) + require.Zero(t, repo.bulkUpdateCalls) +} diff --git a/backend/internal/service/account_grok_media_eligibility_test.go b/backend/internal/service/account_grok_media_eligibility_test.go new file mode 100644 index 000000000..6fbc7e2dd --- /dev/null +++ b/backend/internal/service/account_grok_media_eligibility_test.go @@ -0,0 +1,77 @@ +package service + +import ( + "net/http" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/pkg/xai" + "github.com/stretchr/testify/require" +) + +func TestGrokMediaGenerationEligibility(t *testing.T) { + weeklyUsagePercent := 12.5 + paidBilling := &xai.BillingSummary{ + UsagePercent: &weeklyUsagePercent, + StatusCode: http.StatusOK, + WeeklyStatusCode: http.StatusOK, + } + freeBilling := &xai.BillingSummary{ + StatusCode: http.StatusOK, + WeeklyStatusCode: http.StatusOK, + MonthlyStatusCode: http.StatusOK, + MonthlyUpdatedAt: "2026-07-17T00:00:00Z", + } + inconclusiveBilling := &xai.BillingSummary{ + StatusCode: http.StatusOK, + WeeklyStatusCode: http.StatusOK, + MonthlyStatusCode: http.StatusBadGateway, + Partial: true, + FailedWindows: []string{"monthly"}, + } + + tests := []struct { + name string + account *Account + want bool + wantReason string + }{ + {name: "nil account", account: nil, wantReason: "not_grok"}, + {name: "non grok", account: &Account{Platform: PlatformOpenAI}, wantReason: "not_grok"}, + {name: "non oauth", account: &Account{Platform: PlatformGrok, Type: AccountTypeAPIKey}, want: true, wantReason: "non_oauth"}, + {name: "oauth unobserved", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth}, wantReason: "billing_unobserved"}, + {name: "paid evidence", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: paidBilling}}, want: true, wantReason: "eligible"}, + {name: "free tier", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: freeBilling}}, wantReason: "billing_free_tier"}, + {name: "inconclusive", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: inconclusiveBilling}}, wantReason: "billing_inconclusive"}, + {name: "aggregate forbidden", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: &xai.BillingSummary{StatusCode: http.StatusForbidden}}}, wantReason: "billing_forbidden"}, + {name: "weekly forbidden", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: &xai.BillingSummary{StatusCode: http.StatusOK, WeeklyStatusCode: http.StatusForbidden, MonthlyStatusCode: http.StatusOK}}}, wantReason: "billing_forbidden"}, + {name: "monthly forbidden", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: &xai.BillingSummary{StatusCode: http.StatusOK, WeeklyStatusCode: http.StatusOK, MonthlyStatusCode: http.StatusForbidden}}}, wantReason: "billing_forbidden"}, + {name: "malformed snapshot", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: make(chan int)}}, wantReason: "billing_unobserved"}, + {name: "malformed override ignored", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{GrokMediaEligibleExtraKey: "true", grokBillingExtraKey: paidBilling}}, want: true, wantReason: "eligible"}, + {name: "override disabled", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{GrokMediaEligibleExtraKey: false}}, wantReason: "override_disabled"}, + {name: "override enabled", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{GrokMediaEligibleExtraKey: true, grokBillingExtraKey: &xai.BillingSummary{StatusCode: http.StatusForbidden}}}, want: true, wantReason: "override_enabled"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + eligible, reason := tt.account.GrokMediaGenerationEligibility() + require.Equal(t, tt.want, eligible) + require.Equal(t, tt.wantReason, reason) + }) + } +} + +func TestGrokMediaEndpointCapabilityKeepsOnlyUnobservedOAuthAsProbeCandidate(t *testing.T) { + unobserved := &Account{Platform: PlatformGrok, Type: AccountTypeOAuth} + require.True(t, unobserved.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityGrokMediaGeneration)) + + inconclusive := &Account{ + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Extra: map[string]any{grokBillingExtraKey: &xai.BillingSummary{ + StatusCode: http.StatusOK, + Partial: true, + }}, + } + require.False(t, inconclusive.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityGrokMediaGeneration)) + require.False(t, (&Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth}).SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityGrokMediaGeneration)) +} diff --git a/backend/internal/service/account_service.go b/backend/internal/service/account_service.go index a794d779f..cb4349814 100644 --- a/backend/internal/service/account_service.go +++ b/backend/internal/service/account_service.go @@ -21,30 +21,33 @@ import ( ) var ( - ErrAccountNotFound = infraerrors.NotFound("ACCOUNT_NOT_FOUND", "account not found") - ErrAccountNilInput = infraerrors.BadRequest("ACCOUNT_NIL_INPUT", "account input cannot be nil") - ErrAccountPlatformUnsupported = infraerrors.BadRequest("ACCOUNT_PLATFORM_UNSUPPORTED", "account platform is not supported") - ErrCodexQuotaLimitPercentInvalid = infraerrors.BadRequest("CODEX_QUOTA_LIMIT_PERCENT_INVALID", "Codex quota limit percent must be between 1 and 100") - ErrOwnedAccountAlreadyExists = infraerrors.Conflict("OWNED_ACCOUNT_ALREADY_EXISTS", "account already exists") - ErrOwnedAccountTypeNotAllowed = infraerrors.BadRequest("OWNED_ACCOUNT_TYPE_NOT_ALLOWED", "user accounts only support official OAuth accounts") - ErrOwnedAccountCredentialsInvalid = infraerrors.BadRequest("OWNED_ACCOUNT_CREDENTIALS_INVALID", "OAuth account credentials must include an access token") - ErrOwnedAccountCredentialsNotAllowed = infraerrors.BadRequest("OWNED_ACCOUNT_CREDENTIALS_NOT_ALLOWED", "user accounts cannot include API keys, custom URLs, upstream endpoints, cookies or manual session credentials") - ErrOwnedAccountConcurrencyOutOfRange = infraerrors.BadRequest("OWNED_ACCOUNT_CONCURRENCY_OUT_OF_RANGE", "personal account concurrency must be between 3 and 50") - ErrOwnedAccountLoadFactorOutOfRange = infraerrors.BadRequest("OWNED_ACCOUNT_LOAD_FACTOR_OUT_OF_RANGE", fmt.Sprintf("personal account load factor must be between 1 and %d", AccountMaxLoadFactor)) - ErrOwnedAccountLoadFactorCreditsUnavailable = infraerrors.InternalServer("OWNED_ACCOUNT_LOAD_FACTOR_CREDITS_UNAVAILABLE", "load factor credit accounting is unavailable") - ErrOwnedAccountLoadFactorCreditsInsufficient = infraerrors.BadRequest("OWNED_ACCOUNT_LOAD_FACTOR_CREDITS_INSUFFICIENT", "load factor credits are insufficient") - ErrOwnedAccountLevelNotAllowed = infraerrors.BadRequest("OWNED_ACCOUNT_LEVEL_NOT_ALLOWED", "user accounts cannot manually change account level") - ErrOwnedOpenAIAccountLevelRequired = infraerrors.BadRequest("OWNED_OPENAI_ACCOUNT_LEVEL_REQUIRED", "OpenAI user accounts must select an account level before import") - ErrOwnedAccountProxyRequired = infraerrors.BadRequest("OWNED_ACCOUNT_PROXY_REQUIRED", "user OAuth accounts must use account login with a selected proxy IP") - ErrOwnedOpenAIAccountProxyRequired = ErrOwnedAccountProxyRequired - ErrOwnedAccountGroupPlatformMismatch = infraerrors.BadRequest("OWNED_ACCOUNT_GROUP_PLATFORM_MISMATCH", "account group platform does not match account platform") - ErrOwnedAccountGroupValidationUnavailable = infraerrors.InternalServer("OWNED_ACCOUNT_GROUP_VALIDATION_UNAVAILABLE", "owned account group validation is unavailable") - ErrOwnedAccountPublicPoolUnavailable = infraerrors.BadRequest("OWNED_ACCOUNT_PUBLIC_POOL_UNAVAILABLE", "public shared account pool group is not configured for this account platform") - ErrOwnedAccountPublicPolicyUnavailable = infraerrors.BadRequest("OWNED_ACCOUNT_PUBLIC_POLICY_UNAVAILABLE", "account share policy is not configured for this public account pool") - ErrOwnedAccountPublicValidationFailed = infraerrors.BadRequest("OWNED_ACCOUNT_PUBLIC_VALIDATION_FAILED", "public account validation failed") - ErrOwnedAccountShareModeOnly = infraerrors.BadRequest("OWNED_ACCOUNT_SHARE_MODE_ONLY", "account share mode accounts cannot be moved to the public shared account pool") - ErrOwnedAccountShareModeBoundaryUnavailable = infraerrors.InternalServer("OWNED_ACCOUNT_SHARE_MODE_BOUNDARY_UNAVAILABLE", "account share mode boundary check is unavailable") - ErrOwnedAccountProxyValidationUnavailable = infraerrors.InternalServer("OWNED_ACCOUNT_PROXY_VALIDATION_UNAVAILABLE", "owned account proxy validation is unavailable") + ErrAccountNotFound = infraerrors.NotFound("ACCOUNT_NOT_FOUND", "account not found") + ErrAccountNilInput = infraerrors.BadRequest("ACCOUNT_NIL_INPUT", "account input cannot be nil") + ErrAccountPlatformUnsupported = infraerrors.BadRequest("ACCOUNT_PLATFORM_UNSUPPORTED", "account platform is not supported") + ErrCodexQuotaLimitPercentInvalid = infraerrors.BadRequest("CODEX_QUOTA_LIMIT_PERCENT_INVALID", "Codex quota limit percent must be between 1 and 100") + ErrOwnedAccountAlreadyExists = infraerrors.Conflict("OWNED_ACCOUNT_ALREADY_EXISTS", "account already exists") + ErrOwnedAccountTypeNotAllowed = infraerrors.BadRequest("OWNED_ACCOUNT_TYPE_NOT_ALLOWED", "user accounts only support official OAuth accounts") + ErrOwnedAccountCredentialsInvalid = infraerrors.BadRequest("OWNED_ACCOUNT_CREDENTIALS_INVALID", "OAuth account credentials must include an access token") + ErrOwnedAccountCredentialsNotAllowed = infraerrors.BadRequest("OWNED_ACCOUNT_CREDENTIALS_NOT_ALLOWED", "user accounts cannot include API keys, custom URLs, upstream endpoints, cookies or manual session credentials") + ErrOwnedAgentIdentityCredentialsInvalid = infraerrors.BadRequest("OWNED_AGENT_IDENTITY_CREDENTIALS_INVALID", "Codex Agent Identity credentials are invalid") + ErrOwnedAccountConcurrencyOutOfRange = infraerrors.BadRequest("OWNED_ACCOUNT_CONCURRENCY_OUT_OF_RANGE", "personal account concurrency must be between 3 and 50") + ErrOwnedAccountLoadFactorOutOfRange = infraerrors.BadRequest("OWNED_ACCOUNT_LOAD_FACTOR_OUT_OF_RANGE", fmt.Sprintf("personal account load factor must be between 1 and %d", AccountMaxLoadFactor)) + ErrOwnedAccountLoadFactorCreditsUnavailable = infraerrors.InternalServer("OWNED_ACCOUNT_LOAD_FACTOR_CREDITS_UNAVAILABLE", "load factor credit accounting is unavailable") + ErrOwnedAccountLoadFactorCreditsInsufficient = infraerrors.BadRequest("OWNED_ACCOUNT_LOAD_FACTOR_CREDITS_INSUFFICIENT", "load factor credits are insufficient") + ErrOwnedAccountLevelNotAllowed = infraerrors.BadRequest("OWNED_ACCOUNT_LEVEL_NOT_ALLOWED", "user accounts cannot manually change account level") + ErrOwnedOpenAIAccountLevelRequired = infraerrors.BadRequest("OWNED_OPENAI_ACCOUNT_LEVEL_REQUIRED", "OpenAI user accounts must select an account level before import") + ErrOwnedAccountProxyRequired = infraerrors.BadRequest("OWNED_ACCOUNT_PROXY_REQUIRED", "user OAuth accounts must use account login with a selected proxy IP") + ErrOwnedOpenAIAccountProxyRequired = ErrOwnedAccountProxyRequired + ErrOwnedAccountGroupPlatformMismatch = infraerrors.BadRequest("OWNED_ACCOUNT_GROUP_PLATFORM_MISMATCH", "account group platform does not match account platform") + ErrOwnedAccountGroupValidationUnavailable = infraerrors.InternalServer("OWNED_ACCOUNT_GROUP_VALIDATION_UNAVAILABLE", "owned account group validation is unavailable") + ErrOwnedAccountPublicPoolUnavailable = infraerrors.BadRequest("OWNED_ACCOUNT_PUBLIC_POOL_UNAVAILABLE", "public shared account pool group is not configured for this account platform") + ErrOwnedAccountPublicPolicyUnavailable = infraerrors.BadRequest("OWNED_ACCOUNT_PUBLIC_POLICY_UNAVAILABLE", "account share policy is not configured for this public account pool") + ErrOwnedAccountPublicValidationFailed = infraerrors.BadRequest("OWNED_ACCOUNT_PUBLIC_VALIDATION_FAILED", "public account validation failed") + ErrOwnedAccountShareModeOnly = infraerrors.BadRequest("OWNED_ACCOUNT_SHARE_MODE_ONLY", "account share mode accounts cannot be moved to the public shared account pool") + ErrOwnedAgentIdentityLookupUnavailable = infraerrors.InternalServer("OWNED_AGENT_IDENTITY_LOOKUP_UNAVAILABLE", "Codex Agent Identity account lookup is unavailable") + ErrOwnedAgentIdentityWSInvalidatorUnavailable = infraerrors.InternalServer("OWNED_AGENT_IDENTITY_WS_INVALIDATOR_UNAVAILABLE", "Codex Agent Identity connection invalidation is unavailable") + ErrOwnedAccountShareModeBoundaryUnavailable = infraerrors.InternalServer("OWNED_ACCOUNT_SHARE_MODE_BOUNDARY_UNAVAILABLE", "account share mode boundary check is unavailable") + ErrOwnedAccountProxyValidationUnavailable = infraerrors.InternalServer("OWNED_ACCOUNT_PROXY_VALIDATION_UNAVAILABLE", "owned account proxy validation is unavailable") ) const AccountListGroupUngrouped int64 = -1 @@ -192,19 +195,25 @@ type OwnedPublicShareApprovalOptions struct { AllowRateLimited bool } +type OwnedAccountImportResult struct { + Account *Account + Updated bool +} + // AccountService 账号管理服务 type AccountService struct { - accountRepo AccountRepository - groupRepo GroupRepository - userRepo accountUserRepository - userSubRepo accountSubscriptionLookupRepository - accountSharePolicyRepo AccountSharePolicyRepository - accountShareModeGroups accountShareModeGroupClassifier - settingService *SettingService - privateGroupProvisioner UserPrivateGroupProvisioner - systemNoticeService *SystemNoticeService - proxyRepo ownedAccountProxyRepository - quotaPoolDashboardCache accountQuotaPoolDashboardCache + accountRepo AccountRepository + groupRepo GroupRepository + userRepo accountUserRepository + userSubRepo accountSubscriptionLookupRepository + accountSharePolicyRepo AccountSharePolicyRepository + accountShareModeGroups accountShareModeGroupClassifier + settingService *SettingService + privateGroupProvisioner UserPrivateGroupProvisioner + systemNoticeService *SystemNoticeService + proxyRepo ownedAccountProxyRepository + agentIdentityWSInvalidator agentIdentityWSConnectionInvalidator + quotaPoolDashboardCache accountQuotaPoolDashboardCache } type accountQuotaPoolDashboardCache struct { @@ -247,6 +256,10 @@ type ownedAccountIDBatchRepository interface { ListOwnedAccountIDs(ctx context.Context, ownerUserID int64, accountIDs []int64) ([]int64, error) } +type ownedOpenAIAgentIdentityRepository interface { + GetOwnedOpenAIAgentIdentityByChatGPTAccountID(ctx context.Context, ownerUserID int64, chatGPTAccountID string) (*Account, error) +} + type ownedLoadFactorCreditAccountRepository interface { UpdateOwnedAccountWithLoadFactorCredits(ctx context.Context, ownerUserID int64, account *Account) (*Account, error) } @@ -344,6 +357,13 @@ func (s *AccountService) SetSystemNoticeService(noticeService *SystemNoticeServi s.systemNoticeService = noticeService } +func (s *AccountService) SetAgentIdentityWSInvalidator(invalidator agentIdentityWSConnectionInvalidator) { + if s == nil { + return + } + s.agentIdentityWSInvalidator = invalidator +} + func (s *AccountService) openAIAccountLevelConfigs(ctx context.Context) ([]OpenAIAccountLevelConfig, error) { if s == nil || s.settingService == nil { return DefaultOpenAIAccountLevelConfigs(), nil @@ -523,11 +543,177 @@ func (s *AccountService) EnsureOwnedAccountCanEnterPublicShare(ctx context.Conte } func (s *AccountService) CreateOwned(ctx context.Context, ownerUserID int64, req CreateAccountRequest) (*Account, error) { + if err := rejectOwnedAccountGrokManagedExtra(req.Extra); err != nil { + return nil, err + } return s.createOwned(ctx, ownerUserID, req) } func (s *AccountService) ImportOwned(ctx context.Context, ownerUserID int64, req CreateAccountRequest) (*Account, error) { - return s.createOwned(ctx, ownerUserID, req) + result, err := s.ImportOwnedWithResult(ctx, ownerUserID, req) + if err != nil { + return nil, err + } + return result.Account, nil +} + +func (s *AccountService) ImportOwnedWithResult(ctx context.Context, ownerUserID int64, req CreateAccountRequest) (*OwnedAccountImportResult, error) { + if err := rejectOwnedAccountGrokManagedExtra(req.Extra); err != nil { + return nil, err + } + if !IsOpenAIAgentIdentityCredentials(req.Credentials) { + account, err := s.createOwned(ctx, ownerUserID, req) + if err != nil { + return nil, err + } + return &OwnedAccountImportResult{Account: account}, nil + } + + req.Credentials = normalizeOwnedAgentIdentityCredentials(req.Credentials) + if err := validateOwnedAccountSourceForPlatform(req.Platform, req.Type, req.Credentials, req.Extra); err != nil { + return nil, err + } + chatGPTAccountID := importStringField(req.Credentials, "chatgpt_account_id") + repo, ok := s.accountRepo.(ownedOpenAIAgentIdentityRepository) + if !ok { + return nil, ErrOwnedAgentIdentityLookupUnavailable + } + + existing, err := repo.GetOwnedOpenAIAgentIdentityByChatGPTAccountID(ctx, ownerUserID, chatGPTAccountID) + if err != nil { + return nil, fmt.Errorf("lookup owned Agent Identity account: %w", err) + } + if existing != nil { + account, err := s.updateOwnedAgentIdentityImport(ctx, ownerUserID, existing, req) + if err != nil { + return nil, err + } + return &OwnedAccountImportResult{Account: account, Updated: true}, nil + } + + account, err := s.createOwned(ctx, ownerUserID, req) + if err == nil { + return &OwnedAccountImportResult{Account: account}, nil + } + if !errors.Is(err, ErrOwnedAccountAlreadyExists) { + return nil, err + } + + // A concurrent import may have committed the same owner+Team identity after + // the lookup above. The database unique index is the authority; after that + // conflict becomes visible, converge on the committed row and update it. + existing, lookupErr := repo.GetOwnedOpenAIAgentIdentityByChatGPTAccountID(ctx, ownerUserID, chatGPTAccountID) + if lookupErr != nil { + return nil, fmt.Errorf("reload concurrently imported Agent Identity account: %w", lookupErr) + } + if existing == nil { + return nil, err + } + account, updateErr := s.updateOwnedAgentIdentityImport(ctx, ownerUserID, existing, req) + if updateErr != nil { + return nil, updateErr + } + return &OwnedAccountImportResult{Account: account, Updated: true}, nil +} + +func (s *AccountService) updateOwnedAgentIdentityImport( + ctx context.Context, + ownerUserID int64, + account *Account, + req CreateAccountRequest, +) (*Account, error) { + if account == nil { + return nil, ErrAccountNotFound + } + if account.OwnerUserID == nil || *account.OwnerUserID != ownerUserID || !account.IsOpenAIAgentIdentity() { + return nil, ErrAccountNotFound + } + if s.agentIdentityWSInvalidator == nil { + return nil, ErrOwnedAgentIdentityWSInvalidatorUnavailable + } + if err := validateOwnedAccountSourceForPlatform(req.Platform, req.Type, req.Credentials, req.Extra); err != nil { + return nil, err + } + + before := cloneAccountForNotice(account) + // Rebuild from the Agent Identity allowlist instead of carrying the entire + // historical credential payload forward. Older records may predate the + // recursive credential guard and can contain stale OAuth tokens or other + // fields that Agent Identity must never retain. + nextCredentials := make(map[string]any) + previousRuntimeID := strings.TrimSpace(account.GetCredential("agent_runtime_id")) + nextRuntimeID := importStringField(req.Credentials, "agent_runtime_id") + + allowedCredentialKeys := []string{ + "auth_mode", + "agent_runtime_id", + "agent_private_key", + "task_id", + "chatgpt_account_id", + "chatgpt_user_id", + "email", + "plan_type", + "chatgpt_account_is_fedramp", + } + for _, key := range allowedCredentialKeys { + if value, exists := account.Credentials[key]; exists { + nextCredentials[key] = value + } + } + for _, key := range allowedCredentialKeys { + if key == "task_id" { + continue + } + if value, exists := req.Credentials[key]; exists { + nextCredentials[key] = value + } + } + if taskID := importStringField(req.Credentials, "task_id"); taskID != "" { + nextCredentials["task_id"] = taskID + } else if previousRuntimeID != nextRuntimeID { + delete(nextCredentials, "task_id") + } + if err := validateOwnedAccountSourceForPlatform(PlatformOpenAI, AccountTypeOAuth, nextCredentials, nil); err != nil { + return nil, err + } + + levelConfigs, err := s.openAIAccountLevelConfigs(ctx) + if err != nil { + return nil, err + } + accountLevel := InferOpenAIAccountLevelWithConfigs(nextCredentials, account.Extra, levelConfigs) + if OpenAIAccountLevelConfigByKey(levelConfigs, accountLevel) == nil { + accountLevel = AccountLevelFree + } + + account.Credentials = nextCredentials + account.AccountLevel = accountLevel + var groupIDs []int64 + if NormalizeAccountShareMode(before.ShareMode) == AccountShareModePublic { + // Re-import replaces authentication material. Keep the owner's explicit + // public-share intent, but remove the account from the public pool until + // the handler has completed a fresh connectivity check and approval. + groupIDs, err = s.prepareOwnedPublicShareRevalidation(ctx, ownerUserID, account) + } else { + account.ShareMode = AccountShareModePrivate + account.ShareStatus = AccountShareStatusApproved + account.ErrorMessage = "" + groupIDs, err = s.initialOwnedAccountGroupIDs(ctx, ownerUserID, PlatformOpenAI, AccountTypeOAuth, AccountShareModePrivate, nil) + } + if err != nil { + return nil, err + } + account.ExpiresAt = nil + account.GroupIDs = append([]int64(nil), groupIDs...) + if err := s.accountRepo.Update(ctx, account); err != nil { + return nil, fmt.Errorf("update owned Agent Identity account: %w", err) + } + s.agentIdentityWSInvalidator.InvalidateAgentIdentityWSConnections(account.ID) + if err := s.accountRepo.BindGroups(ctx, account.ID, groupIDs); err != nil { + return nil, fmt.Errorf("bind private Agent Identity account group: %w", err) + } + s.notifyAccountChanged(ctx, before, account) + return account, nil } func (s *AccountService) EnsureOwnedProxyAvailableForNewAccount(ctx context.Context, ownerUserID, proxyID int64) error { @@ -546,17 +732,34 @@ func (s *AccountService) createOwned(ctx context.Context, ownerUserID int64, req if !IsSupportedAccountPlatform(req.Platform) { return nil, ErrAccountPlatformUnsupported } + isAgentIdentity := IsOpenAIAgentIdentityCredentials(req.Credentials) + if isAgentIdentity { + req.Credentials = normalizeOwnedAgentIdentityCredentials(req.Credentials) + } targetLevel := NormalizeAccountLevel(req.AccountLevel) levelConfigs, err := s.openAIAccountLevelConfigs(ctx) if err != nil { return nil, err } - preserveProxy := RequiresUserAccountOAuthProxyWithConfigs(req.Platform, targetLevel, levelConfigs) + preserveProxy := !isAgentIdentity && RequiresUserAccountOAuthProxyWithConfigs(req.Platform, targetLevel, levelConfigs) proxyID := req.ProxyID if err := applyOwnedPersonalAccountTemplateToCreate(&req); err != nil { return nil, err } - if req.Platform == PlatformOpenAI { + if err := validateOwnedAccountSourceForPlatform(req.Platform, req.Type, req.Credentials, req.Extra); err != nil { + return nil, err + } + if isAgentIdentity { + inferredLevel := InferOpenAIAccountLevelWithConfigs(req.Credentials, req.Extra, levelConfigs) + if OpenAIAccountLevelConfigByKey(levelConfigs, inferredLevel) == nil { + inferredLevel = AccountLevelFree + } + targetLevel = inferredLevel + req.AccountLevel = inferredLevel + req.ShareMode = AccountShareModePrivate + req.ProxyID = nil + req.ExpiresAt = nil + } else if req.Platform == PlatformOpenAI { if !IsUserSelectableOpenAIAccountLevelWithConfigs(targetLevel, levelConfigs) { return nil, ErrOwnedOpenAIAccountLevelRequired } @@ -576,9 +779,6 @@ func (s *AccountService) createOwned(ctx context.Context, ownerUserID int64, req } req.AccountLevel = AccountLevelUnknown } - if err := validateOwnedAccountSource(req.Type, req.Credentials, req.Extra); err != nil { - return nil, err - } extra, err := NormalizeCodexQuotaLimitExtra(req.Platform, req.Type, req.Extra) if err != nil { return nil, err @@ -663,20 +863,101 @@ func isAllowedOwnedAccountType(accountType string) bool { return normalized == AccountTypeOAuth } +func normalizeOwnedAgentIdentityCredentials(credentials map[string]any) map[string]any { + normalized := mergeAccountMap(credentials, nil) + normalized["auth_mode"] = OpenAIAuthModeAgentIdentity + for _, key := range []string{ + "agent_runtime_id", + "agent_private_key", + "task_id", + "chatgpt_account_id", + "chatgpt_user_id", + "email", + "plan_type", + } { + if value, ok := normalized[key].(string); ok { + normalized[key] = strings.TrimSpace(value) + } + } + return normalized +} + func validateOwnedAccountSource(accountType string, credentials, extra map[string]any) error { + return validateOwnedAccountSourceForPlatform("", accountType, credentials, extra) +} + +func validateOwnedAccountSourceForPlatform(platform, accountType string, credentials, extra map[string]any) error { if !isAllowedOwnedAccountType(accountType) { return ErrOwnedAccountTypeNotAllowed } - if !hasNonEmptyStringField(credentials, "access_token") { - return ErrOwnedAccountCredentialsInvalid + isAgentIdentity := IsOpenAIAgentIdentityCredentials(credentials) + if isAgentIdentity { + if platform != "" && platform != PlatformOpenAI { + return ErrOwnedAgentIdentityCredentialsInvalid.WithMetadata(map[string]string{"field": "platform"}) + } + for _, key := range []string{"agent_runtime_id", "agent_private_key", "chatgpt_account_id"} { + if !hasNonEmptyStringField(credentials, key) { + return ErrOwnedAgentIdentityCredentialsInvalid.WithMetadata(map[string]string{"field": key}) + } + } + for _, identifier := range []struct { + field string + label string + required bool + }{ + {field: "agent_runtime_id", label: "runtime id", required: true}, + {field: "chatgpt_account_id", label: "Team id", required: true}, + {field: "task_id", label: "task id"}, + {field: "chatgpt_user_id", label: "user id"}, + } { + raw, exists := credentials[identifier.field] + if !exists { + continue + } + value, ok := raw.(string) + if !ok { + return ErrOwnedAgentIdentityCredentialsInvalid.WithMetadata(map[string]string{"field": identifier.field}) + } + if strings.TrimSpace(value) == "" && !identifier.required { + credentials[identifier.field] = "" + continue + } + normalized, err := normalizeAgentIdentityIdentifier(identifier.label, value) + if err != nil { + return ErrOwnedAgentIdentityCredentialsInvalid. + WithMetadata(map[string]string{"field": identifier.field}). + WithCause(err) + } + credentials[identifier.field] = normalized + } + if err := ValidateOpenAIAgentIdentityPrivateKey(importStringField(credentials, "agent_private_key")); err != nil { + return ErrOwnedAgentIdentityCredentialsInvalid.WithMetadata(map[string]string{"field": "agent_private_key"}).WithCause(err) + } + safetyCredentials := mergeAccountMap(credentials, nil) + removeImportMapField(safetyCredentials, "auth_mode") + removeImportMapField(safetyCredentials, "authMode") + if field, ok := findDisallowedOwnedAgentIdentityField(safetyCredentials); ok { + return ErrOwnedAccountCredentialsNotAllowed.WithMetadata(map[string]string{ + "section": "credentials", + "field": field, + }) + } + } else { + if !hasNonEmptyStringField(credentials, "access_token") { + return ErrOwnedAccountCredentialsInvalid + } + if field, ok := findDisallowedOwnedAccountField(credentials); ok { + return ErrOwnedAccountCredentialsNotAllowed.WithMetadata(map[string]string{ + "section": "credentials", + "field": field, + }) + } } - if field, ok := findDisallowedOwnedAccountField(credentials); ok { - return ErrOwnedAccountCredentialsNotAllowed.WithMetadata(map[string]string{ - "section": "credentials", - "field": field, - }) + extraSafetyCheck := findDisallowedOwnedAccountField + if isAgentIdentity { + extraSafetyCheck = findDisallowedOwnedAgentIdentityField } - if field, ok := findDisallowedOwnedAccountField(extra); ok { + if field, ok := extraSafetyCheck(extra); ok { return ErrOwnedAccountCredentialsNotAllowed.WithMetadata(map[string]string{ "section": "extra", "field": field, @@ -742,6 +1023,14 @@ func findDisallowedOwnedAccountField(values map[string]any) (string, bool) { }) } +func findDisallowedOwnedAgentIdentityField(values map[string]any) (string, bool) { + return findDisallowedCredentialContent(values, credentialSafetyOptions{ + AllowOAuthTokenValues: true, + AllowOAuthMetadataURLs: true, + DisallowOAuthTokenFields: true, + }) +} + func normalizeLoadFactor(value *int) *int { if value == nil || *value <= 0 { return nil @@ -918,6 +1207,9 @@ func sanitizeOwnedPersonalAccountUpdate(account *Account, req *UpdateAccountRequ if nextExtra == nil { nextExtra = map[string]any{} } + if err := preserveOwnedAccountGrokManagedExtra(account.Extra, nextExtra); err != nil { + return err + } preserveOwnedPersonalExtraPolicy(account, nextExtra) req.Extra = &nextExtra } @@ -928,6 +1220,9 @@ func ownedPersonalAccountRequiresProxy(account *Account, levelConfigs []OpenAIAc if account == nil { return false } + if account.IsOpenAIAgentIdentity() { + return false + } return RequiresUserAccountOAuthProxyWithConfigs(account.Platform, account.AccountLevel, levelConfigs) } @@ -1299,7 +1594,7 @@ func (s *AccountService) UpdateOwned(ctx context.Context, ownerUserID, accountID groupIDs = managedGroupIDs shouldBindGroups = true } - if err := validateOwnedAccountSource(account.Type, account.Credentials, account.Extra); err != nil { + if err := validateOwnedAccountSourceForPlatform(account.Platform, account.Type, account.Credentials, account.Extra); err != nil { return nil, err } if req.Credentials != nil || req.Extra != nil { @@ -1307,6 +1602,19 @@ func (s *AccountService) UpdateOwned(ctx context.Context, ownerUserID, accountID return nil, err } } + agentIdentityAuthChanged := ownedAgentIdentityAuthMaterialChanged(before, account) + agentIdentityPublicAccessRevoked := ownedAgentIdentityPublicAccessRevoked(before, account) + shouldInvalidateAgentIdentityWS := agentIdentityAuthChanged || agentIdentityPublicAccessRevoked + if shouldInvalidateAgentIdentityWS && s.agentIdentityWSInvalidator == nil { + return nil, ErrOwnedAgentIdentityWSInvalidatorUnavailable + } + if agentIdentityAuthChanged && NormalizeAccountShareMode(account.ShareMode) == AccountShareModePublic { + groupIDs, err = s.prepareOwnedPublicShareRevalidation(ctx, ownerUserID, account) + if err != nil { + return nil, err + } + shouldBindGroups = true + } if !shouldBindGroups && req.GroupIDs != nil { return nil, ErrGroupNotAllowed @@ -1334,6 +1642,9 @@ func (s *AccountService) UpdateOwned(ctx context.Context, ownerUserID, accountID } else if err := s.accountRepo.Update(ctx, account); err != nil { return nil, fmt.Errorf("update account: %w", err) } + if shouldInvalidateAgentIdentityWS { + s.agentIdentityWSInvalidator.InvalidateAgentIdentityWSConnections(account.ID) + } if shouldBindGroups { if err := s.accountRepo.BindGroups(ctx, account.ID, groupIDs); err != nil { return nil, fmt.Errorf("bind groups: %w", err) @@ -1344,6 +1655,34 @@ func (s *AccountService) UpdateOwned(ctx context.Context, ownerUserID, accountID return account, nil } +func ownedAgentIdentityAuthMaterialChanged(before, after *Account) bool { + beforeIsAgentIdentity := before != nil && before.IsOpenAIAgentIdentity() + afterIsAgentIdentity := after != nil && after.IsOpenAIAgentIdentity() + if !beforeIsAgentIdentity && !afterIsAgentIdentity { + return false + } + if beforeIsAgentIdentity != afterIsAgentIdentity { + return true + } + for _, key := range []string{ + "auth_mode", + "agent_runtime_id", + "agent_private_key", + "task_id", + "chatgpt_account_id", + "chatgpt_user_id", + } { + if strings.TrimSpace(before.GetCredential(key)) != strings.TrimSpace(after.GetCredential(key)) { + return true + } + } + return false +} + +func ownedAgentIdentityPublicAccessRevoked(before, after *Account) bool { + return before != nil && after != nil && before.IsOpenAIAgentIdentity() && before.IsPublicShareApproved() && !after.IsPublicShareApproved() +} + func (s *AccountService) SetOwnedOpenAIAccountLevel(ctx context.Context, ownerUserID, accountID int64, accountLevel, reason string) (*Account, error) { account, err := s.GetOwnedByID(ctx, ownerUserID, accountID) if err != nil { @@ -1361,7 +1700,7 @@ func (s *AccountService) SetOwnedOpenAIAccountLevel(ctx context.Context, ownerUs if !IsUserSelectableOpenAIAccountLevelWithConfigs(level, levelConfigs) { return nil, infraerrors.BadRequest("OWNED_ACCOUNT_LEVEL_INVALID", "invalid OpenAI account level") } - if err := validateOwnedAccountSource(account.Type, account.Credentials, account.Extra); err != nil { + if err := validateOwnedAccountSourceForPlatform(account.Platform, account.Type, account.Credentials, account.Extra); err != nil { return nil, err } @@ -1403,10 +1742,17 @@ func (s *AccountService) SetOwnedOpenAIAccountLevel(ctx context.Context, ownerUs return nil, err } } + shouldInvalidateAgentIdentityWS := ownedAgentIdentityPublicAccessRevoked(before, account) + if shouldInvalidateAgentIdentityWS && s.agentIdentityWSInvalidator == nil { + return nil, ErrOwnedAgentIdentityWSInvalidatorUnavailable + } if err := s.accountRepo.Update(ctx, account); err != nil { return nil, fmt.Errorf("update owned OpenAI account level: %w", err) } + if shouldInvalidateAgentIdentityWS { + s.agentIdentityWSInvalidator.InvalidateAgentIdentityWSConnections(account.ID) + } if shouldBindGroups { if err := s.accountRepo.BindGroups(ctx, account.ID, groupIDs); err != nil { return nil, fmt.Errorf("bind account groups after level update: %w", err) @@ -1534,6 +1880,10 @@ func accountDuplicateIdentityKeys(account *Account) []ownedAccountDuplicateKey { if account.Type != AccountTypeOAuth { return nil } + if account.IsOpenAIAgentIdentity() { + add("openai.agent_identity_team", account.GetChatGPTAccountID()) + return keys + } orgID := strings.ToLower(strings.TrimSpace(account.GetOpenAIOrganizationID())) chatGPTUserID := account.GetChatGPTUserID() chatGPTAccountID := account.GetChatGPTAccountID() @@ -1678,6 +2028,9 @@ func (s *AccountService) BulkUpdateOwned(ctx context.Context, ownerUserID int64, if input == nil { return nil, ErrAccountNilInput } + if err := rejectOwnedAccountGrokManagedExtra(input.Extra); err != nil { + return nil, err + } accountIDs := normalizeOwnedBulkAccountIDs(input.AccountIDs) result := &BulkUpdateAccountsResult{ @@ -1755,7 +2108,7 @@ func (s *AccountService) BulkUpdateOwned(ctx context.Context, ownerUserID int64, nextAccount := *account nextAccount.Credentials = nextCredentials nextAccount.Extra = nextExtra - if err := validateOwnedAccountSource(account.Type, nextCredentials, nextExtra); err != nil { + if err := validateOwnedAccountSourceForPlatform(account.Platform, account.Type, nextCredentials, nextExtra); err != nil { return nil, err } nextConcurrency := normalizeOwnedPersonalAccountConcurrency(account.Concurrency) @@ -1975,6 +2328,20 @@ func (s *AccountService) managedOwnedAccountGroupIDsForShareMode(ctx context.Con return s.initialOwnedAccountGroupIDs(ctx, ownerUserID, account.Platform, account.Type, nextMode, nil) } +func (s *AccountService) prepareOwnedPublicShareRevalidation(ctx context.Context, ownerUserID int64, account *Account) ([]int64, error) { + if account == nil { + return nil, ErrAccountNotFound + } + groupIDs, err := s.initialOwnedAccountGroupIDs(ctx, ownerUserID, account.Platform, account.Type, AccountShareModePublic, nil) + if err != nil { + return nil, err + } + account.ShareMode = AccountShareModePublic + account.ShareStatus = AccountShareStatusPending + account.ErrorMessage = "" + return groupIDs, nil +} + func (s *AccountService) ensureAccountCanEnterPublicShare(ctx context.Context, account *Account) error { if account == nil { return ErrAccountNotFound @@ -2006,7 +2373,7 @@ func (s *AccountService) ApproveOwnedPublicShareWithOptions(ctx context.Context, return nil, err } before := cloneAccountForNotice(account) - if err := validateOwnedAccountSource(account.Type, account.Credentials, account.Extra); err != nil { + if err := validateOwnedAccountSourceForPlatform(account.Platform, account.Type, account.Credentials, account.Extra); err != nil { return nil, err } if err := s.ensureAccountCanEnterPublicShare(ctx, account); err != nil { @@ -2030,16 +2397,16 @@ func (s *AccountService) ApproveOwnedPublicShareWithOptions(ctx context.Context, return nil, err } + if err := s.accountRepo.BindGroups(ctx, account.ID, groupIDs); err != nil { + return nil, fmt.Errorf("bind public account groups: %w", err) + } + account.GroupIDs = append([]int64(nil), groupIDs...) account.ShareMode = AccountShareModePublic account.ShareStatus = AccountShareStatusApproved account.ErrorMessage = "" if err := s.accountRepo.Update(ctx, account); err != nil { return nil, fmt.Errorf("update account public share status: %w", err) } - if err := s.accountRepo.BindGroups(ctx, account.ID, groupIDs); err != nil { - return nil, fmt.Errorf("bind public account groups: %w", err) - } - account.GroupIDs = append([]int64(nil), groupIDs...) s.notifyAccountChanged(ctx, before, account) return account, nil } @@ -2076,9 +2443,16 @@ func (s *AccountService) MarkOwnedPublicSharePending(ctx context.Context, ownerU account.ShareMode = AccountShareModePublic account.ShareStatus = AccountShareStatusPending account.ErrorMessage = strings.TrimSpace(reason) + shouldInvalidateAgentIdentityWS := ownedAgentIdentityPublicAccessRevoked(before, account) + if shouldInvalidateAgentIdentityWS && s.agentIdentityWSInvalidator == nil { + return nil, ErrOwnedAgentIdentityWSInvalidatorUnavailable + } if err := s.accountRepo.Update(ctx, account); err != nil { return nil, fmt.Errorf("update account public share status: %w", err) } + if shouldInvalidateAgentIdentityWS { + s.agentIdentityWSInvalidator.InvalidateAgentIdentityWSConnections(account.ID) + } if err := s.accountRepo.BindGroups(ctx, account.ID, groupIDs); err != nil { return nil, fmt.Errorf("bind pending account groups: %w", err) } @@ -2117,9 +2491,16 @@ func (s *AccountService) AutoRepairSuspectedOpenAIFreeAccount(ctx context.Contex return nil, false, err } } + shouldInvalidateAgentIdentityWS := ownedAgentIdentityPublicAccessRevoked(before, account) + if shouldInvalidateAgentIdentityWS && s.agentIdentityWSInvalidator == nil { + return nil, false, ErrOwnedAgentIdentityWSInvalidatorUnavailable + } if err := s.accountRepo.Update(ctx, account); err != nil { return nil, false, fmt.Errorf("update account suspected free repair: %w", err) } + if shouldInvalidateAgentIdentityWS { + s.agentIdentityWSInvalidator.InvalidateAgentIdentityWSConnections(account.ID) + } if account.OwnerUserID != nil { if err := s.accountRepo.BindGroups(ctx, account.ID, groupIDs); err != nil { return nil, false, fmt.Errorf("bind repaired account groups: %w", err) diff --git a/backend/internal/service/account_service_owned_agent_identity_test.go b/backend/internal/service/account_service_owned_agent_identity_test.go new file mode 100644 index 000000000..d5d2d134c --- /dev/null +++ b/backend/internal/service/account_service_owned_agent_identity_test.go @@ -0,0 +1,742 @@ +package service + +import ( + "context" + "errors" + "strings" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/pagination" + "github.com/stretchr/testify/require" +) + +const ( + ownedAgentIdentityPrivateGroupID int64 = 9001 + ownedAgentIdentityPlusPublicGroupID int64 = 9101 + ownedAgentIdentityTeamPublicGroupID int64 = 9102 +) + +type ownedAgentIdentityRepoStub struct { + ownedAccountDuplicateRepoStub + accounts map[int64]*Account + nextID int64 + createCount int + updateCount int + conflictOnNextCreate bool + updateErr error + bindGroupsErr error +} + +func newOwnedAgentIdentityRepoStub() *ownedAgentIdentityRepoStub { + return &ownedAgentIdentityRepoStub{ + accounts: map[int64]*Account{}, + nextID: 1, + } +} + +func cloneOwnedAgentIdentityTestAccount(account *Account) *Account { + if account == nil { + return nil + } + clone := *account + clone.Credentials = mergeAccountMap(account.Credentials, nil) + clone.Extra = mergeAccountMap(account.Extra, nil) + clone.GroupIDs = append([]int64(nil), account.GroupIDs...) + if account.OwnerUserID != nil { + ownerUserID := *account.OwnerUserID + clone.OwnerUserID = &ownerUserID + } + if account.ProxyID != nil { + proxyID := *account.ProxyID + clone.ProxyID = &proxyID + } + if account.ExpiresAt != nil { + expiresAt := *account.ExpiresAt + clone.ExpiresAt = &expiresAt + } + return &clone +} + +func (s *ownedAgentIdentityRepoStub) Create(_ context.Context, account *Account) error { + s.createCount++ + if account.ID <= 0 { + account.ID = s.nextID + s.nextID++ + } + s.accounts[account.ID] = cloneOwnedAgentIdentityTestAccount(account) + if s.conflictOnNextCreate { + s.conflictOnNextCreate = false + return ErrOwnedAccountAlreadyExists + } + return nil +} + +func (s *ownedAgentIdentityRepoStub) Update(_ context.Context, account *Account) error { + s.updateCount++ + if s.updateErr != nil { + return s.updateErr + } + updated := cloneOwnedAgentIdentityTestAccount(account) + if existing := s.accounts[account.ID]; existing != nil { + // GroupIDs model the account_groups relation, which AccountRepository.Update + // does not persist. BindGroups is the only operation that changes it. + updated.GroupIDs = append([]int64(nil), existing.GroupIDs...) + } + s.accounts[account.ID] = updated + return nil +} + +func (s *ownedAgentIdentityRepoStub) BindGroups(_ context.Context, accountID int64, groupIDs []int64) error { + if s.bindGroupsErr != nil { + return s.bindGroupsErr + } + account := s.accounts[accountID] + if account != nil { + account.GroupIDs = append([]int64(nil), groupIDs...) + } + return nil +} + +func (s *ownedAgentIdentityRepoStub) GetByID(_ context.Context, id int64) (*Account, error) { + account := s.accounts[id] + if account == nil { + return nil, ErrAccountNotFound + } + return cloneOwnedAgentIdentityTestAccount(account), nil +} + +func (s *ownedAgentIdentityRepoStub) GetByIDs(_ context.Context, ids []int64) ([]*Account, error) { + accounts := make([]*Account, 0, len(ids)) + for _, id := range ids { + if account := s.accounts[id]; account != nil { + accounts = append(accounts, cloneOwnedAgentIdentityTestAccount(account)) + } + } + return accounts, nil +} + +func (s *ownedAgentIdentityRepoStub) ListOwnedWithFilters( + _ context.Context, + ownerUserID int64, + params pagination.PaginationParams, + platform, accountType, _ string, + _ string, + _, _ int64, + _ string, +) ([]Account, *pagination.PaginationResult, error) { + accounts := make([]Account, 0, len(s.accounts)) + for _, account := range s.accounts { + if account.OwnerUserID == nil || *account.OwnerUserID != ownerUserID { + continue + } + if platform != "" && account.Platform != platform { + continue + } + if accountType != "" && account.Type != accountType { + continue + } + accounts = append(accounts, *cloneOwnedAgentIdentityTestAccount(account)) + } + start := params.Offset() + if start >= len(accounts) { + return []Account{}, &pagination.PaginationResult{Total: int64(len(accounts))}, nil + } + end := start + params.Limit() + if end > len(accounts) { + end = len(accounts) + } + return accounts[start:end], &pagination.PaginationResult{Total: int64(len(accounts))}, nil +} + +func (s *ownedAgentIdentityRepoStub) GetOwnedOpenAIAgentIdentityByChatGPTAccountID( + _ context.Context, + ownerUserID int64, + chatGPTAccountID string, +) (*Account, error) { + chatGPTAccountID = strings.TrimSpace(chatGPTAccountID) + for _, account := range s.accounts { + if account.OwnerUserID == nil || *account.OwnerUserID != ownerUserID || !account.IsOpenAIAgentIdentity() { + continue + } + if strings.TrimSpace(account.GetChatGPTAccountID()) == chatGPTAccountID { + return cloneOwnedAgentIdentityTestAccount(account), nil + } + } + return nil, nil +} + +type recordingAgentIdentityWSInvalidator struct { + accountIDs []int64 +} + +func (r *recordingAgentIdentityWSInvalidator) InvalidateAgentIdentityWSConnections(accountID int64) { + r.accountIDs = append(r.accountIDs, accountID) +} + +func newOwnedAgentIdentityService(repo *ownedAgentIdentityRepoStub) (*AccountService, *recordingAgentIdentityWSInvalidator) { + invalidator := &recordingAgentIdentityWSInvalidator{} + return &AccountService{ + accountRepo: repo, + groupRepo: &ownedPublicShareGroupRepoStub{ + groups: []Group{ + {ID: ownedAgentIdentityPlusPublicGroupID, Name: "PLUS共享号池", Platform: PlatformOpenAI, Status: StatusActive, Scope: GroupScopePublic, RequiredAccountLevel: AccountLevelPlus}, + {ID: ownedAgentIdentityTeamPublicGroupID, Name: "TEAM共享号池", Platform: PlatformOpenAI, Status: StatusActive, Scope: GroupScopePublic, RequiredAccountLevel: AccountLevelTeam}, + }, + }, + accountSharePolicyRepo: &ownedPublicSharePolicyRepoStub{ + policy: &AccountSharePolicy{ID: 1, OwnerShareRatio: 0.7, Enabled: true}, + }, + privateGroupProvisioner: &ownedPrivateGroupProvisionerStub{ + group: &Group{ID: ownedAgentIdentityPrivateGroupID, Name: "OpenAI private", Platform: PlatformOpenAI, Status: StatusActive}, + }, + agentIdentityWSInvalidator: invalidator, + }, invalidator +} + +func ownedAgentIdentityImportRequest(t *testing.T, teamID, userID, runtimeID, planType string) CreateAccountRequest { + t.Helper() + expiresAt := time.Now().Add(24 * time.Hour) + proxyID := int64(77) + return CreateAccountRequest{ + Name: "new import name", + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + ShareMode: AccountShareModePublic, + ProxyID: &proxyID, + Concurrency: 3, + Priority: 5, + ExpiresAt: &expiresAt, + Credentials: map[string]any{ + "auth_mode": OpenAIAuthModeAgentIdentity, + "agent_runtime_id": runtimeID, + "agent_private_key": testAgentIdentityPrivateKey(t), + "chatgpt_account_id": teamID, + "chatgpt_user_id": userID, + "plan_type": planType, + }, + } +} + +func TestAccountServiceImportOwnedAgentIdentityCreatesPrivateOwnedAccount(t *testing.T) { + repo := newOwnedAgentIdentityRepoStub() + svc, invalidator := newOwnedAgentIdentityService(repo) + req := ownedAgentIdentityImportRequest(t, "team-a", "member-a", "runtime-a", "team") + + result, err := svc.ImportOwnedWithResult(context.Background(), 101, req) + + require.NoError(t, err) + require.False(t, result.Updated) + require.NotNil(t, result.Account.OwnerUserID) + require.EqualValues(t, 101, *result.Account.OwnerUserID) + require.Equal(t, AccountShareModePrivate, result.Account.ShareMode) + require.Equal(t, AccountShareStatusApproved, result.Account.ShareStatus) + require.Equal(t, AccountLevelTeam, result.Account.AccountLevel) + require.Nil(t, result.Account.ProxyID) + require.Nil(t, result.Account.ExpiresAt) + require.Equal(t, []int64{ownedAgentIdentityPrivateGroupID}, result.Account.GroupIDs) + require.Empty(t, invalidator.accountIDs) +} + +func TestAccountServiceImportOwnedAgentIdentityUpdatesSameTeamKeepsPrivateAccountPrivate(t *testing.T) { + repo := newOwnedAgentIdentityRepoStub() + svc, invalidator := newOwnedAgentIdentityService(repo) + created, err := svc.ImportOwnedWithResult(context.Background(), 101, ownedAgentIdentityImportRequest(t, "team-a", "member-a", "runtime-old", "team")) + require.NoError(t, err) + + stored := repo.accounts[created.Account.ID] + stored.Name = "preserved local name" + stored.Concurrency = 11 + stored.Priority = 37 + stored.Credentials["task_id"] = "task-old" + stored.Extra["local_setting"] = "keep" + expiresAt := time.Now().Add(time.Hour) + stored.ExpiresAt = &expiresAt + + updated, err := svc.ImportOwnedWithResult(context.Background(), 101, ownedAgentIdentityImportRequest(t, "team-a", "member-a", "runtime-new", "plus")) + + require.NoError(t, err) + require.True(t, updated.Updated) + require.Len(t, repo.accounts, 1) + require.Equal(t, "preserved local name", updated.Account.Name) + require.Equal(t, 11, updated.Account.Concurrency) + require.Equal(t, 37, updated.Account.Priority) + require.Equal(t, "runtime-new", updated.Account.GetCredential("agent_runtime_id")) + require.Empty(t, updated.Account.GetCredential("task_id"), "runtime rotation without a new task must clear the stale task") + require.Equal(t, "keep", updated.Account.Extra["local_setting"]) + require.Equal(t, AccountShareModePrivate, updated.Account.ShareMode) + require.Equal(t, AccountShareStatusApproved, updated.Account.ShareStatus) + require.Equal(t, AccountLevelPlus, updated.Account.AccountLevel) + require.Nil(t, updated.Account.ExpiresAt) + require.Equal(t, []int64{ownedAgentIdentityPrivateGroupID}, updated.Account.GroupIDs) + require.Equal(t, []int64{created.Account.ID}, invalidator.accountIDs) +} + +func TestAccountServiceImportOwnedAgentIdentityRevalidatesChangedPublicAccountWithoutMakingItPrivate(t *testing.T) { + repo := newOwnedAgentIdentityRepoStub() + svc, invalidator := newOwnedAgentIdentityService(repo) + created, err := svc.ImportOwnedWithResult(context.Background(), 101, ownedAgentIdentityImportRequest(t, "team-a", "member-a", "runtime-old", "team")) + require.NoError(t, err) + + publicMode := AccountShareModePublic + pending, err := svc.UpdateOwned(context.Background(), 101, created.Account.ID, UpdateAccountRequest{ShareMode: &publicMode}) + require.NoError(t, err) + require.Equal(t, AccountShareModePublic, pending.ShareMode) + require.Equal(t, AccountShareStatusPending, pending.ShareStatus) + require.Equal(t, []int64{ownedAgentIdentityPrivateGroupID}, pending.GroupIDs) + + approved, err := svc.ApproveOwnedPublicShare(context.Background(), 101, created.Account.ID) + require.NoError(t, err) + require.Equal(t, AccountShareModePublic, approved.ShareMode) + require.Equal(t, AccountShareStatusApproved, approved.ShareStatus) + require.Equal(t, []int64{ownedAgentIdentityPrivateGroupID, ownedAgentIdentityTeamPublicGroupID}, approved.GroupIDs) + + updated, err := svc.ImportOwnedWithResult(context.Background(), 101, ownedAgentIdentityImportRequest(t, "team-a", "member-a", "runtime-new", "plus")) + + require.NoError(t, err) + require.True(t, updated.Updated) + require.Equal(t, AccountShareModePublic, updated.Account.ShareMode, "re-importing an already public account must not silently make it private") + require.Equal(t, AccountShareStatusPending, updated.Account.ShareStatus, "changed authentication material must be revalidated before returning to the public pool") + require.Equal(t, AccountLevelPlus, updated.Account.AccountLevel) + require.Equal(t, "runtime-new", updated.Account.GetCredential("agent_runtime_id")) + require.Equal(t, []int64{ownedAgentIdentityPrivateGroupID}, updated.Account.GroupIDs, "a public account awaiting revalidation must only remain in its owner's private group") + require.Equal(t, []int64{created.Account.ID}, invalidator.accountIDs) + + reapproved, err := svc.ApproveOwnedPublicShare(context.Background(), 101, created.Account.ID) + require.NoError(t, err) + require.Equal(t, AccountShareModePublic, reapproved.ShareMode) + require.Equal(t, AccountShareStatusApproved, reapproved.ShareStatus) + require.Equal(t, []int64{ownedAgentIdentityPrivateGroupID, ownedAgentIdentityPlusPublicGroupID}, reapproved.GroupIDs) +} + +func TestAccountServiceImportOwnedAgentIdentityDropsHistoricalNonAllowlistedCredentials(t *testing.T) { + repo := newOwnedAgentIdentityRepoStub() + svc, _ := newOwnedAgentIdentityService(repo) + req := ownedAgentIdentityImportRequest(t, "team-a", "member-a", "runtime-a", "team") + created, err := svc.ImportOwnedWithResult(context.Background(), 101, req) + require.NoError(t, err) + + stored := repo.accounts[created.Account.ID] + stored.Credentials["task_id"] = "task-existing" + stored.Credentials["accessToken"] = "legacy-access-token" + stored.Credentials["metadata"] = map[string]any{"id_token": "legacy-nested-token"} + stored.Credentials["base_url"] = "https://legacy.invalid" + + updated, err := svc.ImportOwnedWithResult( + context.Background(), + 101, + ownedAgentIdentityImportRequest(t, "team-a", "member-a", "runtime-a", "plus"), + ) + require.NoError(t, err) + require.True(t, updated.Updated) + require.Equal(t, "task-existing", updated.Account.GetCredential("task_id")) + require.Equal(t, "plus", updated.Account.GetCredential("plan_type")) + for _, field := range []string{"accessToken", "metadata", "base_url"} { + require.NotContains(t, updated.Account.Credentials, field) + } + _, hasOAuthToken := findOAuthTokenCredentialContent(updated.Account.Credentials) + require.False(t, hasOAuthToken) + safetyCredentials := mergeAccountMap(updated.Account.Credentials, nil) + removeImportMapField(safetyCredentials, "auth_mode") + _, hasUnsafeField := findDisallowedOwnedAgentIdentityField(safetyCredentials) + require.False(t, hasUnsafeField) +} + +func TestAccountServiceImportOwnedAgentIdentityIsolatesOwnerAndTeam(t *testing.T) { + repo := newOwnedAgentIdentityRepoStub() + svc, _ := newOwnedAgentIdentityService(repo) + + for _, test := range []struct { + ownerID int64 + teamID string + }{ + {ownerID: 101, teamID: "team-a"}, + {ownerID: 101, teamID: "team-b"}, + {ownerID: 202, teamID: "team-a"}, + } { + result, err := svc.ImportOwnedWithResult( + context.Background(), + test.ownerID, + ownedAgentIdentityImportRequest(t, test.teamID, "same-member", "same-runtime", "team"), + ) + require.NoError(t, err) + require.False(t, result.Updated) + } + + require.Len(t, repo.accounts, 3) +} + +func TestAccountServiceImportOwnedAgentIdentityConvergesAfterUniqueConflict(t *testing.T) { + repo := newOwnedAgentIdentityRepoStub() + repo.conflictOnNextCreate = true + svc, invalidator := newOwnedAgentIdentityService(repo) + + result, err := svc.ImportOwnedWithResult(context.Background(), 101, ownedAgentIdentityImportRequest(t, "team-race", "member-a", "runtime-race", "team")) + + require.NoError(t, err) + require.True(t, result.Updated) + require.Len(t, repo.accounts, 1) + require.Equal(t, []int64{result.Account.ID}, invalidator.accountIDs) +} + +func TestValidateOwnedAgentIdentitySourceRejectsUnsafeOrMalformedCredentials(t *testing.T) { + valid := ownedAgentIdentityImportRequest(t, "team-a", "member-a", "runtime-a", "team").Credentials + overlongIdentifier := strings.Repeat("x", agentIdentityIdentifierMaxBytes+1) + + tests := []struct { + name string + platform string + credentials map[string]any + extra map[string]any + }{ + {name: "wrong platform", platform: PlatformAnthropic, credentials: mergeAccountMap(valid, nil)}, + {name: "missing Team", platform: PlatformOpenAI, credentials: mergeAccountMap(valid, map[string]any{"chatgpt_account_id": ""})}, + {name: "bad private key", platform: PlatformOpenAI, credentials: mergeAccountMap(valid, map[string]any{"agent_private_key": "not-a-key"})}, + {name: "mixed access token", platform: PlatformOpenAI, credentials: mergeAccountMap(valid, map[string]any{"access_token": "oauth-token"})}, + {name: "runtime control character", platform: PlatformOpenAI, credentials: mergeAccountMap(valid, map[string]any{"agent_runtime_id": "runtime\x00bad"})}, + {name: "Team control character", platform: PlatformOpenAI, credentials: mergeAccountMap(valid, map[string]any{"chatgpt_account_id": "team\nbad"})}, + {name: "task id too long", platform: PlatformOpenAI, credentials: mergeAccountMap(valid, map[string]any{"task_id": overlongIdentifier})}, + {name: "user id too long", platform: PlatformOpenAI, credentials: mergeAccountMap(valid, map[string]any{"chatgpt_user_id": overlongIdentifier})}, + {name: "custom URL", platform: PlatformOpenAI, credentials: mergeAccountMap(valid, map[string]any{"base_url": "https://evil.example"})}, + {name: "unsafe extra", platform: PlatformOpenAI, credentials: mergeAccountMap(valid, nil), extra: map[string]any{"proxy_url": "https://evil.example"}}, + {name: "nested OAuth token in extra", platform: PlatformOpenAI, credentials: mergeAccountMap(valid, nil), extra: map[string]any{"metadata": []any{map[string]any{"idToken": "must-reject"}}}}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + err := validateOwnedAccountSourceForPlatform(test.platform, AccountTypeOAuth, test.credentials, test.extra) + require.Error(t, err) + }) + } +} + +func TestAccountServiceOwnedAgentIdentityRejectsNestedOAuthTokensBeforeWrite(t *testing.T) { + t.Run("create nested object", func(t *testing.T) { + repo := newOwnedAgentIdentityRepoStub() + svc, _ := newOwnedAgentIdentityService(repo) + req := ownedAgentIdentityImportRequest(t, "team-create", "member-create", "runtime-create", "team") + req.Credentials["metadata"] = map[string]any{ + " ACCESS_TOKEN ": "must-reject", + } + + account, err := svc.CreateOwned(context.Background(), 101, req) + + require.ErrorIs(t, err, ErrOwnedAccountCredentialsNotAllowed) + require.Nil(t, account, "a rejected credential payload must not be returned to the response DTO") + require.Zero(t, repo.createCount) + require.Zero(t, repo.updateCount) + require.Empty(t, repo.accounts) + }) + + t.Run("import nested array", func(t *testing.T) { + repo := newOwnedAgentIdentityRepoStub() + svc, _ := newOwnedAgentIdentityService(repo) + req := ownedAgentIdentityImportRequest(t, "team-import", "member-import", "runtime-import", "team") + req.Credentials["metadata"] = []any{ + map[string]any{"RefreshToken": "must-reject"}, + } + + result, err := svc.ImportOwnedWithResult(context.Background(), 101, req) + + require.ErrorIs(t, err, ErrOwnedAccountCredentialsNotAllowed) + require.Nil(t, result, "a rejected credential payload must not be returned to the response DTO") + require.Zero(t, repo.createCount) + require.Zero(t, repo.updateCount) + require.Empty(t, repo.accounts) + }) + + t.Run("update nested array object", func(t *testing.T) { + repo := newOwnedAgentIdentityRepoStub() + svc, _ := newOwnedAgentIdentityService(repo) + created, err := svc.ImportOwnedWithResult( + context.Background(), + 101, + ownedAgentIdentityImportRequest(t, "team-update", "member-update", "runtime-update", "team"), + ) + require.NoError(t, err) + require.NotNil(t, created) + + credentials := mergeAccountMap(created.Account.Credentials, map[string]any{ + "metadata": []any{ + map[string]any{" ID_TOKEN ": "must-reject"}, + }, + }) + account, err := svc.UpdateOwned( + context.Background(), + 101, + created.Account.ID, + UpdateAccountRequest{Credentials: &credentials}, + ) + + require.ErrorIs(t, err, ErrOwnedAccountCredentialsNotAllowed) + require.Nil(t, account, "a rejected credential payload must not be returned to the response DTO") + require.Zero(t, repo.updateCount) + require.NotContains(t, repo.accounts[created.Account.ID].Credentials, "metadata") + }) +} + +func TestAccountServiceUpdateOwnedAgentIdentityInvalidatesWSWhenAuthMaterialChanges(t *testing.T) { + repo := newOwnedAgentIdentityRepoStub() + svc, invalidator := newOwnedAgentIdentityService(repo) + created, err := svc.ImportOwnedWithResult(context.Background(), 101, ownedAgentIdentityImportRequest(t, "team-a", "member-a", "runtime-old", "team")) + require.NoError(t, err) + + credentials := map[string]any{ + "auth_mode": OpenAIAuthModeAgentIdentity, + "agent_runtime_id": "runtime-new", + "task_id": "task-new", + "chatgpt_account_id": "team-a", + "chatgpt_user_id": "member-a", + "plan_type": "team", + } + updated, err := svc.UpdateOwned(context.Background(), 101, created.Account.ID, UpdateAccountRequest{Credentials: &credentials}) + + require.NoError(t, err) + require.Equal(t, "runtime-new", updated.GetCredential("agent_runtime_id")) + require.Equal(t, "task-new", updated.GetCredential("task_id")) + require.Equal(t, []int64{created.Account.ID}, invalidator.accountIDs) +} + +func TestAccountServiceUpdateOwnedPublicAgentIdentityFailsClosedWhenPrivateGroupBindingFails(t *testing.T) { + repo := newOwnedAgentIdentityRepoStub() + svc, invalidator := newOwnedAgentIdentityService(repo) + created, err := svc.ImportOwnedWithResult(context.Background(), 101, ownedAgentIdentityImportRequest(t, "team-a", "member-a", "runtime-old", "team")) + require.NoError(t, err) + + publicMode := AccountShareModePublic + _, err = svc.UpdateOwned(context.Background(), 101, created.Account.ID, UpdateAccountRequest{ShareMode: &publicMode}) + require.NoError(t, err) + approved, err := svc.ApproveOwnedPublicShare(context.Background(), 101, created.Account.ID) + require.NoError(t, err) + require.Equal(t, []int64{ownedAgentIdentityPrivateGroupID, ownedAgentIdentityTeamPublicGroupID}, approved.GroupIDs) + + repo.updateCount = 0 + invalidator.accountIDs = nil + repo.bindGroupsErr = errors.New("injected group binding failure") + credentials := mergeAccountMap(approved.Credentials, map[string]any{ + "agent_runtime_id": "runtime-new", + "task_id": "task-new", + }) + + updated, err := svc.UpdateOwned( + context.Background(), + 101, + created.Account.ID, + UpdateAccountRequest{Credentials: &credentials}, + ) + + require.ErrorContains(t, err, "bind groups") + require.Nil(t, updated) + require.Equal(t, 1, repo.updateCount) + require.Equal(t, []int64{created.Account.ID}, invalidator.accountIDs) + stored := repo.accounts[created.Account.ID] + require.Equal(t, "runtime-new", stored.GetCredential("agent_runtime_id")) + require.Equal(t, AccountShareModePublic, stored.ShareMode) + require.Equal(t, AccountShareStatusPending, stored.ShareStatus) + require.Equal(t, []int64{ownedAgentIdentityPrivateGroupID, ownedAgentIdentityTeamPublicGroupID}, stored.GroupIDs) + require.False(t, stored.IsVisibleToConsumer(202), "a stale public-group binding must remain fail-closed after the row enters pending") +} + +func TestAccountServiceUpdateOwnedPublicAgentIdentityRevocationFailsClosedWhenGroupBindingFails(t *testing.T) { + repo := newOwnedAgentIdentityRepoStub() + svc, invalidator := newOwnedAgentIdentityService(repo) + created, err := svc.ImportOwnedWithResult(context.Background(), 101, ownedAgentIdentityImportRequest(t, "team-a", "member-a", "runtime-old", "team")) + require.NoError(t, err) + + publicMode := AccountShareModePublic + _, err = svc.UpdateOwned(context.Background(), 101, created.Account.ID, UpdateAccountRequest{ShareMode: &publicMode}) + require.NoError(t, err) + approved, err := svc.ApproveOwnedPublicShare(context.Background(), 101, created.Account.ID) + require.NoError(t, err) + require.Equal(t, []int64{ownedAgentIdentityPrivateGroupID, ownedAgentIdentityTeamPublicGroupID}, approved.GroupIDs) + + repo.updateCount = 0 + repo.bindGroupsErr = errors.New("injected group binding failure") + privateMode := AccountShareModePrivate + updated, err := svc.UpdateOwned( + context.Background(), + 101, + created.Account.ID, + UpdateAccountRequest{ShareMode: &privateMode}, + ) + + require.ErrorContains(t, err, "bind groups") + require.Nil(t, updated) + require.Equal(t, 1, repo.updateCount) + require.Equal(t, []int64{created.Account.ID}, invalidator.accountIDs) + stored := repo.accounts[created.Account.ID] + require.Equal(t, AccountShareModePrivate, stored.ShareMode) + require.Equal(t, AccountShareStatusApproved, stored.ShareStatus) + require.Equal(t, []int64{ownedAgentIdentityPrivateGroupID, ownedAgentIdentityTeamPublicGroupID}, stored.GroupIDs) + require.False(t, stored.IsVisibleToConsumer(202), "private status must block stale public-group membership") +} + +func TestAccountServiceUpdateOwnedPublicAgentIdentityToPrivateInvalidatesWS(t *testing.T) { + repo := newOwnedAgentIdentityRepoStub() + svc, invalidator := newOwnedAgentIdentityService(repo) + created, err := svc.ImportOwnedWithResult(context.Background(), 101, ownedAgentIdentityImportRequest(t, "team-a", "member-a", "runtime-old", "team")) + require.NoError(t, err) + + publicMode := AccountShareModePublic + _, err = svc.UpdateOwned(context.Background(), 101, created.Account.ID, UpdateAccountRequest{ShareMode: &publicMode}) + require.NoError(t, err) + _, err = svc.ApproveOwnedPublicShare(context.Background(), 101, created.Account.ID) + require.NoError(t, err) + invalidator.accountIDs = nil + + privateMode := AccountShareModePrivate + updated, err := svc.UpdateOwned(context.Background(), 101, created.Account.ID, UpdateAccountRequest{ShareMode: &privateMode}) + + require.NoError(t, err) + require.Equal(t, AccountShareModePrivate, updated.ShareMode) + require.Equal(t, []int64{ownedAgentIdentityPrivateGroupID}, updated.GroupIDs) + require.Equal(t, []int64{created.Account.ID}, invalidator.accountIDs) +} + +func TestAccountServiceSetOwnedAgentIdentityLevelSuspensionInvalidatesWS(t *testing.T) { + repo := newOwnedAgentIdentityRepoStub() + svc, invalidator := newOwnedAgentIdentityService(repo) + created, err := svc.ImportOwnedWithResult(context.Background(), 101, ownedAgentIdentityImportRequest(t, "team-a", "member-a", "runtime-old", "team")) + require.NoError(t, err) + + publicMode := AccountShareModePublic + _, err = svc.UpdateOwned(context.Background(), 101, created.Account.ID, UpdateAccountRequest{ShareMode: &publicMode}) + require.NoError(t, err) + _, err = svc.ApproveOwnedPublicShare(context.Background(), 101, created.Account.ID) + require.NoError(t, err) + invalidator.accountIDs = nil + + updated, err := svc.SetOwnedOpenAIAccountLevel(context.Background(), 101, created.Account.ID, AccountLevelFree, "free tier detected") + + require.NoError(t, err) + require.Equal(t, AccountShareStatusSuspended, updated.ShareStatus) + require.Equal(t, []int64{created.Account.ID}, invalidator.accountIDs) +} + +func TestAccountServiceAutoRepairOwnedAgentIdentitySuspensionInvalidatesWS(t *testing.T) { + repo := newOwnedAgentIdentityRepoStub() + svc, invalidator := newOwnedAgentIdentityService(repo) + created, err := svc.ImportOwnedWithResult(context.Background(), 101, ownedAgentIdentityImportRequest(t, "team-a", "member-a", "runtime-old", "team")) + require.NoError(t, err) + + publicMode := AccountShareModePublic + _, err = svc.UpdateOwned(context.Background(), 101, created.Account.ID, UpdateAccountRequest{ShareMode: &publicMode}) + require.NoError(t, err) + _, err = svc.ApproveOwnedPublicShare(context.Background(), 101, created.Account.ID) + require.NoError(t, err) + groupRepo := svc.groupRepo.(*ownedPublicShareGroupRepoStub) + groupRepo.groups = append(groupRepo.groups, Group{ID: 9100, Name: "FREE共享号池", Platform: PlatformOpenAI, Status: StatusActive, Scope: GroupScopePublic, RequiredAccountLevel: AccountLevelFree}) + account := repo.accounts[created.Account.ID] + now := time.Now().UTC() + account.Extra = mergeAccountMap(account.Extra, map[string]any{ + "quota_weekly_limit": 50.0, + "codex_7d_used_percent": 100.0, + "codex_7d_reset_at": now.Add(24 * time.Hour).Format(time.RFC3339), + }) + invalidator.accountIDs = nil + + updated, repaired, err := svc.AutoRepairSuspectedOpenAIFreeAccount(context.Background(), created.Account.ID, 60, "quota proof") + + require.NoError(t, err) + require.True(t, repaired) + require.Equal(t, AccountShareStatusSuspended, updated.ShareStatus) + require.Equal(t, []int64{created.Account.ID}, invalidator.accountIDs) +} + +func TestAccountServiceUpdateOwnedAgentIdentityFailsBeforeWriteWithoutWSInvalidator(t *testing.T) { + repo := newOwnedAgentIdentityRepoStub() + svc, _ := newOwnedAgentIdentityService(repo) + created, err := svc.ImportOwnedWithResult(context.Background(), 101, ownedAgentIdentityImportRequest(t, "team-a", "member-a", "runtime-old", "team")) + require.NoError(t, err) + svc.agentIdentityWSInvalidator = nil + + credentials := map[string]any{ + "auth_mode": OpenAIAuthModeAgentIdentity, + "agent_runtime_id": "runtime-new", + "chatgpt_account_id": "team-a", + "chatgpt_user_id": "member-a", + } + _, err = svc.UpdateOwned(context.Background(), 101, created.Account.ID, UpdateAccountRequest{Credentials: &credentials}) + + require.ErrorIs(t, err, ErrOwnedAgentIdentityWSInvalidatorUnavailable) + require.Zero(t, repo.updateCount) + require.Equal(t, "runtime-old", repo.accounts[created.Account.ID].GetCredential("agent_runtime_id")) +} + +func TestOwnedAgentIdentityCanEnterPublicShareAfterApproval(t *testing.T) { + repo := newOwnedAgentIdentityRepoStub() + svc, _ := newOwnedAgentIdentityService(repo) + created, err := svc.ImportOwnedWithResult(context.Background(), 101, ownedAgentIdentityImportRequest(t, "team-a", "member-a", "runtime-a", "team")) + require.NoError(t, err) + require.Equal(t, AccountShareModePrivate, created.Account.ShareMode, "Agent Identity imports must default to private") + require.Equal(t, AccountShareStatusApproved, created.Account.ShareStatus) + require.Equal(t, []int64{ownedAgentIdentityPrivateGroupID}, created.Account.GroupIDs) + + publicMode := AccountShareModePublic + pending, err := svc.UpdateOwned(context.Background(), 101, created.Account.ID, UpdateAccountRequest{ShareMode: &publicMode}) + require.NoError(t, err) + require.Equal(t, AccountShareModePublic, pending.ShareMode) + require.Equal(t, AccountShareStatusPending, pending.ShareStatus) + require.Equal(t, []int64{ownedAgentIdentityPrivateGroupID}, pending.GroupIDs, "an unverified account must not enter the public pool") + + approved, err := svc.ApproveOwnedPublicShare(context.Background(), 101, created.Account.ID) + require.NoError(t, err) + require.Equal(t, AccountShareModePublic, approved.ShareMode) + require.Equal(t, AccountShareStatusApproved, approved.ShareStatus) + require.Equal(t, []int64{ownedAgentIdentityPrivateGroupID, ownedAgentIdentityTeamPublicGroupID}, approved.GroupIDs) +} + +func TestAccountServiceApproveOwnedPublicShareKeepsPendingWhenGroupBindingFails(t *testing.T) { + repo := newOwnedAgentIdentityRepoStub() + svc, _ := newOwnedAgentIdentityService(repo) + created, err := svc.ImportOwnedWithResult(context.Background(), 101, ownedAgentIdentityImportRequest(t, "team-a", "member-a", "runtime-a", "team")) + require.NoError(t, err) + + publicMode := AccountShareModePublic + _, err = svc.UpdateOwned(context.Background(), 101, created.Account.ID, UpdateAccountRequest{ShareMode: &publicMode}) + require.NoError(t, err) + repo.updateCount = 0 + repo.bindGroupsErr = errors.New("injected group binding failure") + + approved, err := svc.ApproveOwnedPublicShare(context.Background(), 101, created.Account.ID) + + require.Nil(t, approved) + require.ErrorContains(t, err, "bind public account groups") + require.Zero(t, repo.updateCount, "approval must not be persisted before group binding succeeds") + stored := repo.accounts[created.Account.ID] + require.Equal(t, AccountShareModePublic, stored.ShareMode) + require.Equal(t, AccountShareStatusPending, stored.ShareStatus) + require.Equal(t, []int64{ownedAgentIdentityPrivateGroupID}, stored.GroupIDs) + require.False(t, stored.IsVisibleToConsumer(202)) +} + +func TestAccountServiceApproveOwnedPublicShareRetriesAfterStatusUpdateFailure(t *testing.T) { + repo := newOwnedAgentIdentityRepoStub() + svc, _ := newOwnedAgentIdentityService(repo) + created, err := svc.ImportOwnedWithResult(context.Background(), 101, ownedAgentIdentityImportRequest(t, "team-a", "member-a", "runtime-a", "team")) + require.NoError(t, err) + + publicMode := AccountShareModePublic + _, err = svc.UpdateOwned(context.Background(), 101, created.Account.ID, UpdateAccountRequest{ShareMode: &publicMode}) + require.NoError(t, err) + repo.updateCount = 0 + repo.updateErr = errors.New("injected status update failure") + + approved, err := svc.ApproveOwnedPublicShare(context.Background(), 101, created.Account.ID) + + require.Nil(t, approved) + require.ErrorContains(t, err, "update account public share status") + require.Equal(t, 1, repo.updateCount) + stored := repo.accounts[created.Account.ID] + require.Equal(t, AccountShareModePublic, stored.ShareMode) + require.Equal(t, AccountShareStatusPending, stored.ShareStatus) + require.Equal(t, []int64{ownedAgentIdentityPrivateGroupID, ownedAgentIdentityTeamPublicGroupID}, stored.GroupIDs) + require.False(t, stored.IsVisibleToConsumer(202), "pending status must block a partially bound account") + + repo.updateErr = nil + retried, err := svc.ApproveOwnedPublicShare(context.Background(), 101, created.Account.ID) + require.NoError(t, err) + require.Equal(t, AccountShareStatusApproved, retried.ShareStatus) + require.Equal(t, []int64{ownedAgentIdentityPrivateGroupID, ownedAgentIdentityTeamPublicGroupID}, retried.GroupIDs) + require.True(t, retried.IsVisibleToConsumer(202)) +} diff --git a/backend/internal/service/account_share_mode.go b/backend/internal/service/account_share_mode.go index 970dcd466..42616c753 100644 --- a/backend/internal/service/account_share_mode.go +++ b/backend/internal/service/account_share_mode.go @@ -1614,7 +1614,7 @@ func (s *AccountShareModeService) CreateOpenAIListingFromToken(ctx context.Conte if input.AutoPauseOnExpired != nil { account.AutoPauseOnExpired = *input.AutoPauseOnExpired } - if err := validateOwnedAccountSource(account.Type, account.Credentials, account.Extra); err != nil { + if err := validateOwnedAccountSourceForPlatform(account.Platform, account.Type, account.Credentials, account.Extra); err != nil { return nil, err } listing := &AccountShareListing{ @@ -1716,7 +1716,7 @@ func (s *AccountShareModeService) CreateAnthropicListingFromToken(ctx context.Co if input.AutoPauseOnExpired != nil { account.AutoPauseOnExpired = *input.AutoPauseOnExpired } - if err := validateOwnedAccountSource(account.Type, account.Credentials, account.Extra); err != nil { + if err := validateOwnedAccountSourceForPlatform(account.Platform, account.Type, account.Credentials, account.Extra); err != nil { return nil, err } listing := &AccountShareListing{ diff --git a/backend/internal/service/account_usage_service.go b/backend/internal/service/account_usage_service.go index 0f6a42dd7..7a9ffb3d5 100644 --- a/backend/internal/service/account_usage_service.go +++ b/backend/internal/service/account_usage_service.go @@ -61,9 +61,9 @@ type UsageLogRepository interface { // User dashboard stats GetUserDashboardStats(ctx context.Context, userID int64) (*usagestats.UserDashboardStats, error) GetAPIKeyDashboardStats(ctx context.Context, apiKeyID int64) (*usagestats.UserDashboardStats, error) - GetUserUsageTrendByUserID(ctx context.Context, userID int64, startTime, endTime time.Time, granularity string) ([]usagestats.TrendDataPoint, error) + GetUserUsageTrendByUserID(ctx context.Context, userID int64, startTime, endTime time.Time, granularity string, location *time.Location) ([]usagestats.TrendDataPoint, error) GetUserModelStats(ctx context.Context, userID int64, startTime, endTime time.Time) ([]usagestats.ModelStat, error) - GetUserAccountSharingDashboard(ctx context.Context, userID int64, startTime, endTime time.Time, granularity string) (*usagestats.AccountSharingDashboardStats, error) + GetUserAccountSharingDashboard(ctx context.Context, userID int64, startTime, endTime time.Time, granularity string, location *time.Location) (*usagestats.AccountSharingDashboardStats, error) // Admin usage listing/stats ListWithFilters(ctx context.Context, params pagination.PaginationParams, filters usagestats.UsageLogFilters) ([]UsageLog, *pagination.PaginationResult, error) diff --git a/backend/internal/service/admin_service.go b/backend/internal/service/admin_service.go index 432765dcb..3ceb9ad50 100644 --- a/backend/internal/service/admin_service.go +++ b/backend/internal/service/admin_service.go @@ -569,7 +569,13 @@ const ( proxyQualityClientUserAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/136.0.0.0 Safari/537.36" ) -var ErrRPMStatusUnavailable = infraerrors.New(http.StatusNotImplemented, "RPM_STATUS_UNAVAILABLE", "RPM cache not available") +var ( + ErrRPMStatusUnavailable = infraerrors.New(http.StatusNotImplemented, "RPM_STATUS_UNAVAILABLE", "RPM cache not available") + errAdminBulkOwnedAgentIdentityAuthUpdateUnsupported = infraerrors.BadRequest( + "ACCOUNT_BULK_OWNED_AGENT_IDENTITY_AUTH_UPDATE_UNSUPPORTED", + "Codex Agent Identity authentication material must be updated one account at a time", + ) +) // adminServiceImpl implements AdminService type adminServiceImpl struct { @@ -594,6 +600,7 @@ type adminServiceImpl struct { privacyClientFactory PrivacyClientFactory privateGroupProvisioner UserPrivateGroupProvisioner systemNoticeService *SystemNoticeService + agentIdentityWSInvalidator agentIdentityWSConnectionInvalidator } type userGroupRateBatchReader interface { @@ -659,6 +666,13 @@ func SetAdminSystemNoticeService(svc AdminService, noticeService *SystemNoticeSe return svc } +func SetAdminAgentIdentityWSInvalidator(svc AdminService, invalidator agentIdentityWSConnectionInvalidator) AdminService { + if impl, ok := svc.(*adminServiceImpl); ok { + impl.agentIdentityWSInvalidator = invalidator + } + return svc +} + func (s *adminServiceImpl) openAIAccountLevelConfigs(ctx context.Context) ([]OpenAIAccountLevelConfig, error) { if s == nil || s.settingService == nil { return DefaultOpenAIAccountLevelConfigs(), nil @@ -3230,6 +3244,9 @@ func (s *adminServiceImpl) prepareAccountCreate(ctx context.Context, input *Crea if !IsSupportedAccountPlatform(input.Platform) { return nil, nil, ErrAccountPlatformUnsupported } + if _, err := validateAdminGrokManagedExtra(input.Extra); err != nil { + return nil, nil, err + } extra, err := NormalizeCodexQuotaLimitExtra(input.Platform, input.Type, input.Extra) if err != nil { return nil, nil, err @@ -3392,6 +3409,49 @@ func shouldEnsureOAuthPrivacyAfterCreate(account *Account) bool { return account != nil && account.Type == AccountTypeOAuth && !account.IsOpenAIAgentIdentity() } +func isOwnedOpenAIAgentIdentity(account *Account) bool { + return account != nil && + account.OwnerUserID != nil && + *account.OwnerUserID > 0 && + account.IsOpenAIAgentIdentity() +} + +func agentIdentityOwnerUserIDChanged(before, after *Account) bool { + if (before == nil || !before.IsOpenAIAgentIdentity()) && (after == nil || !after.IsOpenAIAgentIdentity()) { + return false + } + beforeOwnerID := int64(0) + if before != nil && before.OwnerUserID != nil { + beforeOwnerID = *before.OwnerUserID + } + afterOwnerID := int64(0) + if after != nil && after.OwnerUserID != nil { + afterOwnerID = *after.OwnerUserID + } + return beforeOwnerID != afterOwnerID +} + +func shouldForceAdminOwnedAgentIdentityPending( + before *Account, + after *Account, + input *UpdateAccountInput, + authMaterialChanged bool, + ownerUserIDChanged bool, +) bool { + if after == nil || (!isOwnedOpenAIAgentIdentity(before) && !isOwnedOpenAIAgentIdentity(after)) { + return false + } + if NormalizeAccountShareMode(after.ShareMode) != AccountShareModePublic || + NormalizeAccountShareStatus(after.ShareStatus) == AccountShareStatusSuspended { + return false + } + enteredPublic := before == nil || NormalizeAccountShareMode(before.ShareMode) != AccountShareModePublic + explicitlyApproved := input != nil && + strings.TrimSpace(input.ShareStatus) != "" && + NormalizeAccountShareStatus(input.ShareStatus) == AccountShareStatusApproved + return enteredPublic || authMaterialChanged || ownerUserIDChanged || explicitlyApproved +} + func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *UpdateAccountInput) (*Account, error) { account, err := s.accountRepo.GetByID(ctx, id) if err != nil { @@ -3421,6 +3481,14 @@ func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *U // Extra 使用 map：需要区分“未提供(nil)”与“显式清空({})”。 // 关闭配额限制时前端会删除 quota_* 键并提交 extra:{}，此时也必须落库。 if input.Extra != nil { + mediaOverrideProvided, err := validateAdminGrokManagedExtra(input.Extra) + if err != nil { + return nil, err + } + if !mediaOverrideProvided { + preserveMapKey(account.Extra, input.Extra, GrokMediaEligibleExtraKey) + } + preserveMapKey(account.Extra, input.Extra, grokBillingExtraKey) // 保留配额用量字段，防止编辑账号时意外重置 for _, key := range []string{"quota_used", "quota_daily_used", "quota_daily_start", "quota_weekly_used", "quota_weekly_start"} { if v, ok := account.Extra[key]; ok { @@ -3577,9 +3645,25 @@ func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *U } } + agentIdentityAuthChanged := ownedAgentIdentityAuthMaterialChanged(before, account) + agentIdentityOwnerChanged := agentIdentityOwnerUserIDChanged(before, account) + if shouldForceAdminOwnedAgentIdentityPending(before, account, input, agentIdentityAuthChanged, agentIdentityOwnerChanged) { + account.ShareStatus = AccountShareStatusPending + account.ErrorMessage = "" + } + shouldInvalidateAgentIdentityWS := agentIdentityAuthChanged || + agentIdentityOwnerChanged || + ownedAgentIdentityPublicAccessRevoked(before, account) + if shouldInvalidateAgentIdentityWS && s.agentIdentityWSInvalidator == nil { + return nil, ErrOwnedAgentIdentityWSInvalidatorUnavailable + } + if err := s.accountRepo.Update(ctx, account); err != nil { return nil, err } + if shouldInvalidateAgentIdentityWS { + s.agentIdentityWSInvalidator.InvalidateAgentIdentityWSConnections(account.ID) + } // 绑定分组 if input.GroupIDs != nil { @@ -3603,6 +3687,9 @@ func (s *adminServiceImpl) BulkUpdateAccounts(ctx context.Context, input *BulkUp if input == nil { return nil, infraerrors.BadRequest("ACCOUNT_BULK_UPDATE_INVALID", "bulk update input is required") } + if _, err := validateAdminGrokManagedExtra(input.Extra); err != nil { + return nil, err + } if len(input.AccountIDs) == 0 && input.Filters != nil { accountIDs, err := s.resolveBulkUpdateTargetIDs(ctx, input.Filters) @@ -3644,6 +3731,31 @@ func (s *adminServiceImpl) BulkUpdateAccounts(ctx context.Context, input *BulkUp return preflightAccounts, nil } + var agentIdentityWSInvalidationIDs []int64 + if len(input.Credentials) > 0 { + accounts, err := loadPreflightAccounts() + if err != nil { + return nil, err + } + for _, account := range accounts { + if account == nil { + continue + } + after := cloneAccountForNotice(account) + after.Credentials = mergeAccountMapPreservingSensitiveCreds(account.Credentials, input.Credentials) + if !ownedAgentIdentityAuthMaterialChanged(account, after) { + continue + } + if isOwnedOpenAIAgentIdentity(account) || isOwnedOpenAIAgentIdentity(after) { + return nil, errAdminBulkOwnedAgentIdentityAuthUpdateUnsupported + } + agentIdentityWSInvalidationIDs = append(agentIdentityWSInvalidationIDs, account.ID) + } + if len(agentIdentityWSInvalidationIDs) > 0 && s.agentIdentityWSInvalidator == nil { + return nil, ErrOwnedAgentIdentityWSInvalidatorUnavailable + } + } + if input.GroupIDs != nil || input.AccountLevel != nil || len(input.Credentials) > 0 || len(input.Extra) > 0 { accounts, err := loadPreflightAccounts() if err != nil { @@ -3844,6 +3956,9 @@ func (s *adminServiceImpl) BulkUpdateAccounts(ctx context.Context, input *BulkUp if _, err := s.accountRepo.BulkUpdate(ctx, input.AccountIDs, repoUpdates); err != nil { return nil, err } + for _, accountID := range agentIdentityWSInvalidationIDs { + s.agentIdentityWSInvalidator.InvalidateAgentIdentityWSConnections(accountID) + } // Handle group bindings per account (requires individual operations). for _, accountID := range input.AccountIDs { diff --git a/backend/internal/service/admin_service_bulk_update_test.go b/backend/internal/service/admin_service_bulk_update_test.go index eee5977e2..3c70a6687 100644 --- a/backend/internal/service/admin_service_bulk_update_test.go +++ b/backend/internal/service/admin_service_bulk_update_test.go @@ -790,3 +790,143 @@ func TestAdminService_BulkUpdateAccounts_RejectsHigherOpenAILevelIntoLowerPool(t require.Empty(t, repo.bulkUpdateIDs) require.Empty(t, repo.boundGroupIDs) } + +func newAdminOwnedAgentIdentityTestAccount(t *testing.T, ownerUserID int64, shareMode, shareStatus, runtimeID string) *Account { + t.Helper() + return &Account{ + ID: 1, + Name: "owned-agent-identity", + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + AccountLevel: AccountLevelTeam, + Credentials: map[string]any{ + "auth_mode": OpenAIAuthModeAgentIdentity, + "agent_runtime_id": runtimeID, + "agent_private_key": testAgentIdentityPrivateKey(t), + "task_id": "task-old", + "chatgpt_account_id": "team-a", + "chatgpt_user_id": "member-a", + }, + OwnerUserID: &ownerUserID, + ShareMode: shareMode, + ShareStatus: shareStatus, + Concurrency: 3, + Status: StatusActive, + Schedulable: true, + } +} + +func TestAdminServiceUpdateOwnedPublicAgentIdentityAuthChangeForcesPendingAndInvalidatesWS(t *testing.T) { + account := newAdminOwnedAgentIdentityTestAccount(t, 101, AccountShareModePublic, AccountShareStatusApproved, "runtime-old") + repo := &accountRepoStubForBulkUpdate{getByIDAccounts: map[int64]*Account{account.ID: account}} + invalidator := &recordingAgentIdentityWSInvalidator{} + svc := &adminServiceImpl{accountRepo: repo, agentIdentityWSInvalidator: invalidator} + credentials := mergeAccountMap(account.Credentials, map[string]any{"agent_runtime_id": "runtime-new"}) + + updated, err := svc.UpdateAccount(context.Background(), account.ID, &UpdateAccountInput{Credentials: credentials}) + + require.NoError(t, err) + require.Equal(t, AccountShareStatusPending, updated.ShareStatus) + require.Equal(t, "runtime-new", updated.GetCredential("agent_runtime_id")) + require.Equal(t, []int64{account.ID}, invalidator.accountIDs) +} + +func TestAdminServiceUpdateOwnedPublicAgentIdentityToPrivateInvalidatesWS(t *testing.T) { + account := newAdminOwnedAgentIdentityTestAccount(t, 101, AccountShareModePublic, AccountShareStatusApproved, "runtime-old") + repo := &accountRepoStubForBulkUpdate{getByIDAccounts: map[int64]*Account{account.ID: account}} + invalidator := &recordingAgentIdentityWSInvalidator{} + svc := &adminServiceImpl{accountRepo: repo, agentIdentityWSInvalidator: invalidator} + + updated, err := svc.UpdateAccount(context.Background(), account.ID, &UpdateAccountInput{ShareMode: AccountShareModePrivate}) + + require.NoError(t, err) + require.Equal(t, AccountShareModePrivate, updated.ShareMode) + require.Equal(t, AccountShareStatusApproved, updated.ShareStatus) + require.Equal(t, []int64{account.ID}, invalidator.accountIDs) +} + +func TestAdminServiceUpdateOwnedPublicAgentIdentitySuspensionIsPreserved(t *testing.T) { + account := newAdminOwnedAgentIdentityTestAccount(t, 101, AccountShareModePublic, AccountShareStatusApproved, "runtime-old") + repo := &accountRepoStubForBulkUpdate{getByIDAccounts: map[int64]*Account{account.ID: account}} + invalidator := &recordingAgentIdentityWSInvalidator{} + svc := &adminServiceImpl{accountRepo: repo, agentIdentityWSInvalidator: invalidator} + + updated, err := svc.UpdateAccount(context.Background(), account.ID, &UpdateAccountInput{ShareStatus: AccountShareStatusSuspended}) + + require.NoError(t, err) + require.Equal(t, AccountShareStatusSuspended, updated.ShareStatus) + require.Equal(t, []int64{account.ID}, invalidator.accountIDs) +} + +func TestAdminServiceUpdateOwnedPublicAgentIdentityExplicitApprovalForcesPending(t *testing.T) { + account := newAdminOwnedAgentIdentityTestAccount(t, 101, AccountShareModePublic, AccountShareStatusApproved, "runtime-old") + repo := &accountRepoStubForBulkUpdate{getByIDAccounts: map[int64]*Account{account.ID: account}} + invalidator := &recordingAgentIdentityWSInvalidator{} + svc := &adminServiceImpl{accountRepo: repo, agentIdentityWSInvalidator: invalidator} + + updated, err := svc.UpdateAccount(context.Background(), account.ID, &UpdateAccountInput{ShareStatus: AccountShareStatusApproved}) + + require.NoError(t, err) + require.Equal(t, AccountShareStatusPending, updated.ShareStatus) + require.Equal(t, []int64{account.ID}, invalidator.accountIDs) +} + +func TestAdminServiceUpdateOwnedPublicAgentIdentityOwnerChangeForcesPending(t *testing.T) { + account := newAdminOwnedAgentIdentityTestAccount(t, 101, AccountShareModePublic, AccountShareStatusApproved, "runtime-old") + repo := &accountRepoStubForBulkUpdate{getByIDAccounts: map[int64]*Account{account.ID: account}} + invalidator := &recordingAgentIdentityWSInvalidator{} + svc := &adminServiceImpl{accountRepo: repo, agentIdentityWSInvalidator: invalidator} + nextOwnerUserID := int64(202) + + updated, err := svc.UpdateAccount(context.Background(), account.ID, &UpdateAccountInput{OwnerUserID: &nextOwnerUserID}) + + require.NoError(t, err) + require.NotNil(t, updated.OwnerUserID) + require.Equal(t, nextOwnerUserID, *updated.OwnerUserID) + require.Equal(t, AccountShareStatusPending, updated.ShareStatus) + require.Equal(t, []int64{account.ID}, invalidator.accountIDs) +} + +func TestAdminServiceUpdateOwnedAgentIdentityFailsBeforeWriteWithoutWSInvalidator(t *testing.T) { + account := newAdminOwnedAgentIdentityTestAccount(t, 101, AccountShareModePublic, AccountShareStatusApproved, "runtime-old") + repo := &accountRepoStubForBulkUpdate{getByIDAccounts: map[int64]*Account{account.ID: account}} + svc := &adminServiceImpl{accountRepo: repo} + + updated, err := svc.UpdateAccount(context.Background(), account.ID, &UpdateAccountInput{ShareMode: AccountShareModePrivate}) + + require.Nil(t, updated) + require.ErrorIs(t, err, ErrOwnedAgentIdentityWSInvalidatorUnavailable) + require.Nil(t, repo.updatedAccount) +} + +func TestAdminServiceBulkUpdateOwnedAgentIdentityAuthChangeFailsBeforeWrite(t *testing.T) { + account := newAdminOwnedAgentIdentityTestAccount(t, 101, AccountShareModePrivate, AccountShareStatusApproved, "runtime-old") + repo := &accountRepoStubForBulkUpdate{getByIDsAccounts: []*Account{account}} + svc := &adminServiceImpl{accountRepo: repo} + + result, err := svc.BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{ + AccountIDs: []int64{account.ID}, + Credentials: map[string]any{"agent_runtime_id": "runtime-new"}, + }) + + require.Nil(t, result) + require.True(t, infraerrors.IsBadRequest(err)) + require.Equal(t, "ACCOUNT_BULK_OWNED_AGENT_IDENTITY_AUTH_UPDATE_UNSUPPORTED", infraerrors.Reason(err)) + require.Empty(t, repo.bulkUpdateIDs) +} + +func TestAdminServiceBulkUpdateOwnedAgentIdentityNonAuthCredentialsDoesNotFalseReject(t *testing.T) { + account := newAdminOwnedAgentIdentityTestAccount(t, 101, AccountShareModePrivate, AccountShareStatusApproved, "runtime-old") + repo := &accountRepoStubForBulkUpdate{getByIDsAccounts: []*Account{account}} + svc := &adminServiceImpl{accountRepo: repo} + + result, err := svc.BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{ + AccountIDs: []int64{account.ID}, + Credentials: map[string]any{"email": "updated@example.com"}, + }) + + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, []int64{account.ID}, repo.bulkUpdateIDs) + require.Equal(t, "updated@example.com", repo.bulkUpdateUpdate.Credentials["email"]) +} diff --git a/backend/internal/service/api_key_service.go b/backend/internal/service/api_key_service.go index 702bb91fd..230fa0b3c 100644 --- a/backend/internal/service/api_key_service.go +++ b/backend/internal/service/api_key_service.go @@ -10,6 +10,7 @@ import ( "strings" "sync" "time" + "unicode/utf8" "github.com/Wei-Shaw/sub2api/internal/config" infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" @@ -25,11 +26,15 @@ var ( ErrGroupNotAllowed = infraerrors.Forbidden("GROUP_NOT_ALLOWED", "user is not allowed to bind this group") ErrAPIKeyExists = infraerrors.Conflict("API_KEY_EXISTS", "api key already exists") ErrAPIKeyTooShort = infraerrors.BadRequest("API_KEY_TOO_SHORT", "api key must be at least 16 characters") + ErrAPIKeyTooLong = infraerrors.BadRequest("API_KEY_TOO_LONG", "api key must be at most 128 characters") ErrAPIKeyInvalidChars = infraerrors.BadRequest("API_KEY_INVALID_CHARS", "api key can only contain letters, numbers, underscores, and hyphens") ErrAPIKeyRateLimited = infraerrors.TooManyRequests("API_KEY_RATE_LIMITED", "too many failed attempts, please try again later") ErrInvalidIPPattern = infraerrors.BadRequest("INVALID_IP_PATTERN", "invalid IP or CIDR pattern") ErrAPIKeyGroupRequired = infraerrors.BadRequest("API_KEY_GROUP_REQUIRED", "api key group is required when ungrouped key scheduling is disabled") ErrAPIKeyGroupRouteInvalid = infraerrors.BadRequest("API_KEY_GROUP_ROUTE_INVALID", "invalid api key group route") + ErrAPIKeyExpirationConflict = infraerrors.BadRequest("API_KEY_EXPIRATION_CONFLICT", "expires_at and expires_in_days cannot be provided together") + ErrAPIKeyExpirationInvalid = infraerrors.BadRequest("API_KEY_EXPIRATION_INVALID", "expires_at must be a valid RFC3339 timestamp") + ErrAPIKeyExpirationNotFuture = infraerrors.BadRequest("API_KEY_EXPIRATION_NOT_FUTURE", "expires_at must be later than the current time") ErrAPIKeyAccountShareBindingExists = infraerrors.Conflict( "API_KEY_ACCOUNT_SHARE_BINDING_EXISTS", "api key is bound to active or queued account share mode usage", @@ -47,8 +52,12 @@ var ( ) const ( - apiKeyMaxErrorsPerHour = 20 - apiKeyLastUsedMinTouch = 30 * time.Second + MaxAPIKeyCredentialCharacters = 128 + // MaxAPIKeyCredentialBytes bounds request work before the character-level + // check while preserving credentials generated with a multi-byte prefix. + MaxAPIKeyCredentialBytes = MaxAPIKeyCredentialCharacters * utf8.UTFMax + apiKeyMaxErrorsPerHour = 20 + apiKeyLastUsedMinTouch = 30 * time.Second // DB 写失败后的短退避，避免请求路径持续同步重试造成写风暴与高延迟。 apiKeyLastUsedFailBackoff = 5 * time.Second ) @@ -173,8 +182,9 @@ type CreateAPIKeyRequest struct { IPBlacklist []string `json:"ip_blacklist"` // IP 黑名单 // Quota fields - Quota float64 `json:"quota"` // Quota limit in USD (0 = unlimited) - ExpiresInDays *int `json:"expires_in_days"` // Days until expiry (nil = never expires) + Quota float64 `json:"quota"` // Quota limit in USD (0 = unlimited) + ExpiresInDays *int `json:"expires_in_days"` // Days until expiry (nil = never expires) + ExpiresAt *time.Time `json:"expires_at"` // Exact expiration time (nil = never expires) // Rate limit fields (0 = unlimited) RateLimit5h float64 `json:"rate_limit_5h"` @@ -188,8 +198,8 @@ type UpdateAPIKeyRequest struct { GroupID *int64 `json:"group_id"` GroupRoutes *[]APIKeyGroupRoute `json:"group_routes"` Status *string `json:"status"` - IPWhitelist []string `json:"ip_whitelist"` // IP 白名单（空数组清空） - IPBlacklist []string `json:"ip_blacklist"` // IP 黑名单（空数组清空） + IPWhitelist *[]string `json:"ip_whitelist"` // IP 白名单（nil 不修改，空数组清空） + IPBlacklist *[]string `json:"ip_blacklist"` // IP 黑名单（nil 不修改，空数组清空） // Quota fields Quota *float64 `json:"quota"` // Quota limit in USD (nil = no change, 0 = unlimited) @@ -314,6 +324,9 @@ func (s *APIKeyService) GenerateKey() (string, error) { } key := prefix + hex.EncodeToString(bytes) + if !apiKeyCredentialWithinLimit(key) { + return "", fmt.Errorf("default.api_key_prefix must produce a valid UTF-8 api key of at most %d characters", MaxAPIKeyCredentialCharacters) + } return key, nil } @@ -323,6 +336,9 @@ func (s *APIKeyService) ValidateCustomKey(key string) error { if len(key) < 16 { return ErrAPIKeyTooShort } + if utf8.RuneCountInString(key) > MaxAPIKeyCredentialCharacters { + return ErrAPIKeyTooLong + } // 检查字符：只允许字母、数字、下划线、连字符 for _, c := range key { @@ -466,8 +482,31 @@ func (s *APIKeyService) validateAPIKeyGroupRoutes(ctx context.Context, user *Use return nil } -// Create 鍒涘缓API Key +func resolveCreateAPIKeyExpiration(req CreateAPIKeyRequest, now time.Time) (*time.Time, error) { + if req.ExpiresAt != nil && req.ExpiresInDays != nil { + return nil, ErrAPIKeyExpirationConflict + } + if req.ExpiresAt != nil { + if !req.ExpiresAt.After(now) { + return nil, ErrAPIKeyExpirationNotFuture + } + expiresAt := *req.ExpiresAt + return &expiresAt, nil + } + if req.ExpiresInDays != nil && *req.ExpiresInDays > 0 { + expiresAt := now.AddDate(0, 0, *req.ExpiresInDays) + return &expiresAt, nil + } + return nil, nil +} + +// Create 创建 API Key。 func (s *APIKeyService) Create(ctx context.Context, userID int64, req CreateAPIKeyRequest) (*APIKey, error) { + expiresAt, err := resolveCreateAPIKeyExpiration(req, time.Now()) + if err != nil { + return nil, err + } + // 验证用户存在 user, err := s.userRepo.GetByID(ctx, userID) if err != nil { @@ -571,12 +610,7 @@ func (s *APIKeyService) Create(ctx context.Context, userID int64, req CreateAPIK RateLimit5h: req.RateLimit5h, RateLimit1d: req.RateLimit1d, RateLimit7d: req.RateLimit7d, - } - - // Set expiration time if specified - if req.ExpiresInDays != nil && *req.ExpiresInDays > 0 { - expiresAt := time.Now().AddDate(0, 0, *req.ExpiresInDays) - apiKey.ExpiresAt = &expiresAt + ExpiresAt: expiresAt, } if err := s.apiKeyRepo.Create(ctx, apiKey); err != nil { @@ -656,6 +690,9 @@ func (s *APIKeyService) GetByID(ctx context.Context, id int64) (*APIKey, error) // GetByKey 根据Key字符串获取API Key（用于认证） func (s *APIKeyService) GetByKey(ctx context.Context, key string) (*APIKey, error) { + if !apiKeyCredentialWithinLimit(key) { + return nil, ErrAPIKeyNotFound + } cacheKey := s.authCacheKey(key) if entry, ok := s.getAuthCacheEntry(ctx, cacheKey); ok { @@ -706,6 +743,12 @@ func (s *APIKeyService) GetByKey(ctx context.Context, key string) (*APIKey, erro return apiKey, nil } +func apiKeyCredentialWithinLimit(key string) bool { + return len(key) <= MaxAPIKeyCredentialBytes && + utf8.ValidString(key) && + utf8.RuneCountInString(key) <= MaxAPIKeyCredentialCharacters +} + // Update 更新API Key func (s *APIKeyService) Update(ctx context.Context, id int64, userID int64, req UpdateAPIKeyRequest) (*APIKey, error) { apiKey, err := s.apiKeyRepo.GetByID(ctx, id) @@ -721,15 +764,15 @@ func (s *APIKeyService) Update(ctx context.Context, id int64, userID int64, req currentAPIKey.GroupRoutes = append([]APIKeyGroupRoute(nil), apiKey.GroupRoutes...) // 验证 IP 白名单格式 - if len(req.IPWhitelist) > 0 { - if invalid := ip.ValidateIPPatterns(req.IPWhitelist); len(invalid) > 0 { + if req.IPWhitelist != nil && len(*req.IPWhitelist) > 0 { + if invalid := ip.ValidateIPPatterns(*req.IPWhitelist); len(invalid) > 0 { return nil, fmt.Errorf("%w: %v", ErrInvalidIPPattern, invalid) } } // 验证 IP 黑名单格式 - if len(req.IPBlacklist) > 0 { - if invalid := ip.ValidateIPPatterns(req.IPBlacklist); len(invalid) > 0 { + if req.IPBlacklist != nil && len(*req.IPBlacklist) > 0 { + if invalid := ip.ValidateIPPatterns(*req.IPBlacklist); len(invalid) > 0 { return nil, fmt.Errorf("%w: %v", ErrInvalidIPPattern, invalid) } } @@ -822,9 +865,13 @@ func (s *APIKeyService) Update(ctx context.Context, id int64, userID int64, req } } - // 更新 IP 限制（空数组会清空设置） - apiKey.IPWhitelist = req.IPWhitelist - apiKey.IPBlacklist = req.IPBlacklist + // 更新 IP 限制（nil 不修改，空数组清空设置） + if req.IPWhitelist != nil { + apiKey.IPWhitelist = *req.IPWhitelist + } + if req.IPBlacklist != nil { + apiKey.IPBlacklist = *req.IPBlacklist + } // Update rate limit configuration if req.RateLimit5h != nil { diff --git a/backend/internal/service/api_key_service_create_expiration_test.go b/backend/internal/service/api_key_service_create_expiration_test.go new file mode 100644 index 000000000..08e488ecd --- /dev/null +++ b/backend/internal/service/api_key_service_create_expiration_test.go @@ -0,0 +1,118 @@ +package service + +import ( + "context" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/stretchr/testify/require" +) + +type apiKeyCreateRepoStub struct { + APIKeyRepository + created *APIKey +} + +func (s *apiKeyCreateRepoStub) Create(_ context.Context, key *APIKey) error { + clone := *key + if key.ExpiresAt != nil { + expiresAt := *key.ExpiresAt + clone.ExpiresAt = &expiresAt + } + s.created = &clone + return nil +} + +type apiKeyCreateUserRepoStub struct { + UserRepository + user *User +} + +func (s *apiKeyCreateUserRepoStub) GetByID(context.Context, int64) (*User, error) { + clone := *s.user + return &clone, nil +} + +type apiKeyCreateGroupRepoStub struct { + GroupRepository + group *Group +} + +func (s *apiKeyCreateGroupRepoStub) GetByID(context.Context, int64) (*Group, error) { + clone := *s.group + return &clone, nil +} + +func TestAPIKeyServiceCreatePreservesExactExpiration(t *testing.T) { + groupID := int64(9) + expiresAt := time.Date(2099, time.March, 3, 21, 6, 7, 123456789, time.UTC) + repo := &apiKeyCreateRepoStub{} + svc := &APIKeyService{ + apiKeyRepo: repo, + userRepo: &apiKeyCreateUserRepoStub{user: &User{ID: 42}}, + groupRepo: &apiKeyCreateGroupRepoStub{group: &Group{ + ID: groupID, + Status: StatusActive, + Scope: GroupScopePublic, + SubscriptionType: SubscriptionTypeStandard, + }}, + cfg: &config.Config{}, + } + + created, err := svc.Create(context.Background(), 42, CreateAPIKeyRequest{ + Name: "precise expiration", + GroupID: &groupID, + ExpiresAt: &expiresAt, + }) + + require.NoError(t, err) + require.NotNil(t, repo.created) + require.NotNil(t, created.ExpiresAt) + require.Equal(t, expiresAt, *created.ExpiresAt) + require.Equal(t, expiresAt, *repo.created.ExpiresAt) +} + +func TestAPIKeyServiceCreateRejectsInvalidExactExpirationBeforeRepositoryAccess(t *testing.T) { + now := time.Now() + past := now.Add(-time.Minute) + legacyDays := 30 + + tests := []struct { + name string + req CreateAPIKeyRequest + want error + }{ + { + name: "past expires_at", + req: CreateAPIKeyRequest{ExpiresAt: &past}, + want: ErrAPIKeyExpirationNotFuture, + }, + { + name: "conflicting expiration fields", + req: CreateAPIKeyRequest{ExpiresAt: &past, ExpiresInDays: &legacyDays}, + want: ErrAPIKeyExpirationConflict, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + svc := &APIKeyService{} + + _, err := svc.Create(context.Background(), 42, tt.req) + + require.ErrorIs(t, err, tt.want) + }) + } +} + +func TestResolveCreateAPIKeyExpirationKeepsLegacyDaysCompatibility(t *testing.T) { + now := time.Date(2026, time.July, 23, 8, 9, 10, 11, time.UTC) + days := 30 + + expiresAt, err := resolveCreateAPIKeyExpiration(CreateAPIKeyRequest{ExpiresInDays: &days}, now) + + require.NoError(t, err) + require.NotNil(t, expiresAt) + require.Equal(t, now.AddDate(0, 0, days), *expiresAt) +} diff --git a/backend/internal/service/api_key_service_delete_test.go b/backend/internal/service/api_key_service_delete_test.go index abb3fe07a..2e3a29877 100644 --- a/backend/internal/service/api_key_service_delete_test.go +++ b/backend/internal/service/api_key_service_delete_test.go @@ -328,6 +328,55 @@ func TestAPIKeyServiceUpdateAllowsUnrelatedChangeWithAccountShareBinding(t *test require.Len(t, repo.updatedKeys, 1) } +func TestAPIKeyServiceUpdatePreservesIPRestrictionsWhenOmitted(t *testing.T) { + repo := &apiKeyRepoStub{apiKey: &APIKey{ + ID: 42, + UserID: 7, + Key: "k", + Status: StatusAPIKeyActive, + IPWhitelist: []string{"192.0.2.10", "198.51.100.0/24"}, + IPBlacklist: []string{"203.0.113.9"}, + }} + svc := &APIKeyService{apiKeyRepo: repo} + + updated, err := svc.Update(context.Background(), 42, 7, UpdateAPIKeyRequest{}) + require.NoError(t, err) + require.Equal(t, []string{"192.0.2.10", "198.51.100.0/24"}, updated.IPWhitelist) + require.Equal(t, []string{"203.0.113.9"}, updated.IPBlacklist) + require.Len(t, repo.updatedKeys, 1) + require.Equal(t, updated.IPWhitelist, repo.updatedKeys[0].IPWhitelist) + require.Equal(t, updated.IPBlacklist, repo.updatedKeys[0].IPBlacklist) +} + +func TestAPIKeyServiceUpdateClearsIPRestrictionsWhenExplicitlyEmpty(t *testing.T) { + repo := &apiKeyRepoStub{apiKey: &APIKey{ + ID: 42, + UserID: 7, + Key: "k", + Status: StatusAPIKeyActive, + IPWhitelist: []string{"192.0.2.10"}, + IPBlacklist: []string{"203.0.113.9"}, + }} + svc := &APIKeyService{apiKeyRepo: repo} + emptyWhitelist := []string{} + emptyBlacklist := []string{} + + updated, err := svc.Update(context.Background(), 42, 7, UpdateAPIKeyRequest{ + IPWhitelist: &emptyWhitelist, + IPBlacklist: &emptyBlacklist, + }) + require.NoError(t, err) + require.NotNil(t, updated.IPWhitelist) + require.Empty(t, updated.IPWhitelist) + require.NotNil(t, updated.IPBlacklist) + require.Empty(t, updated.IPBlacklist) + require.Len(t, repo.updatedKeys, 1) + require.NotNil(t, repo.updatedKeys[0].IPWhitelist) + require.Empty(t, repo.updatedKeys[0].IPWhitelist) + require.NotNil(t, repo.updatedKeys[0].IPBlacklist) + require.Empty(t, repo.updatedKeys[0].IPBlacklist) +} + func TestAPIKeyServiceUpdateDetectsInPlaceGroupRouteChanges(t *testing.T) { groupID := int64(10) routes := []APIKeyGroupRoute{{GroupID: groupID, Priority: 100, Weight: 2, Enabled: true, CooldownSeconds: 30}} diff --git a/backend/internal/service/api_key_service_length_test.go b/backend/internal/service/api_key_service_length_test.go new file mode 100644 index 000000000..405242a0f --- /dev/null +++ b/backend/internal/service/api_key_service_length_test.go @@ -0,0 +1,46 @@ +package service + +import ( + "context" + "strings" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/stretchr/testify/require" +) + +func TestValidateCustomKeyRejectsMoreThanMaximumCharacters(t *testing.T) { + service := &APIKeyService{} + + require.NoError(t, service.ValidateCustomKey(strings.Repeat("a", MaxAPIKeyCredentialCharacters))) + require.ErrorIs(t, service.ValidateCustomKey(strings.Repeat("a", MaxAPIKeyCredentialCharacters+1)), ErrAPIKeyTooLong) +} + +func TestAPIKeyCredentialWithinLimitPreservesMultiBytePrefixes(t *testing.T) { + require.True(t, apiKeyCredentialWithinLimit(strings.Repeat("界", MaxAPIKeyCredentialCharacters))) + require.False(t, apiKeyCredentialWithinLimit(strings.Repeat("界", MaxAPIKeyCredentialCharacters+1))) + require.False(t, apiKeyCredentialWithinLimit(string([]byte{0xff}))) +} + +func TestGetByKeyRejectsTooManyCharactersBeforeRepositoryAccess(t *testing.T) { + service := &APIKeyService{} + + apiKey, err := service.GetByKey(context.Background(), strings.Repeat("a", MaxAPIKeyCredentialCharacters+1)) + + require.Nil(t, apiKey) + require.ErrorIs(t, err, ErrAPIKeyNotFound) +} + +func TestGenerateKeyValidatesConfiguredPrefixLength(t *testing.T) { + service := &APIKeyService{cfg: &config.Config{Default: config.DefaultConfig{ + APIKeyPrefix: strings.Repeat("界", MaxAPIKeyCredentialCharacters-64), + }}} + + key, err := service.GenerateKey() + require.NoError(t, err) + require.Equal(t, MaxAPIKeyCredentialCharacters, len([]rune(key))) + + service.cfg.Default.APIKeyPrefix += "界" + _, err = service.GenerateKey() + require.ErrorContains(t, err, "default.api_key_prefix") +} diff --git a/backend/internal/service/backup_service.go b/backend/internal/service/backup_service.go index a9c2c6bd3..6be40dac1 100644 --- a/backend/internal/service/backup_service.go +++ b/backend/internal/service/backup_service.go @@ -26,17 +26,27 @@ const ( settingKeyBackupSchedule = "backup_schedule" settingKeyBackupRecords = "backup_records" settingKeyUsageRetention = "usage_cleanup_auto_retention" - - maxBackupRecords = 100 + backupEncryptedSecretV1 = "enc:v1:" + + // 在 OPS 04:00 清理前的低峰时段执行；标准 cron 表达式不会在服务启动时立即触发。 + backupExpirationCleanupCronExpr = "30 3 * * *" + backupExpirationCleanupBatchSize = 50 + backupExpirationCleanupAttempts = 3 + backupExpirationCleanupRetryWait = 10 * time.Minute + backupCompensationTimeout = 30 * time.Second ) var ( - ErrBackupS3NotConfigured = infraerrors.BadRequest("BACKUP_S3_NOT_CONFIGURED", "backup S3 storage is not configured") - ErrBackupNotFound = infraerrors.NotFound("BACKUP_NOT_FOUND", "backup record not found") - ErrBackupInProgress = infraerrors.Conflict("BACKUP_IN_PROGRESS", "a backup is already in progress") - ErrRestoreInProgress = infraerrors.Conflict("RESTORE_IN_PROGRESS", "a restore is already in progress") - ErrBackupRecordsCorrupt = infraerrors.InternalServer("BACKUP_RECORDS_CORRUPT", "backup records data is corrupted") - ErrBackupS3ConfigCorrupt = infraerrors.InternalServer("BACKUP_S3_CONFIG_CORRUPT", "backup S3 config data is corrupted") + ErrBackupS3NotConfigured = infraerrors.BadRequest("BACKUP_S3_NOT_CONFIGURED", "backup S3 storage is not configured") + ErrBackupNotFound = infraerrors.NotFound("BACKUP_NOT_FOUND", "backup record not found") + ErrBackupInProgress = infraerrors.Conflict("BACKUP_IN_PROGRESS", "a backup is already in progress") + ErrRestoreInProgress = infraerrors.Conflict("RESTORE_IN_PROGRESS", "a restore is already in progress") + ErrBackupRecordsCorrupt = infraerrors.InternalServer("BACKUP_RECORDS_CORRUPT", "backup records data is corrupted") + ErrBackupS3ConfigCorrupt = infraerrors.InternalServer("BACKUP_S3_CONFIG_CORRUPT", "backup S3 config data is corrupted") + ErrSecretEncryptionKeyNotConfigured = infraerrors.BadRequest( + "SECRET_ENCRYPTION_KEY_NOT_CONFIGURED", + "cannot store the S3 secret access key: configure a fixed TOTP_ENCRYPTION_KEY before saving durable credentials", + ) ) // ─── 接口定义 ─── @@ -134,8 +144,10 @@ type BackupService struct { dbCfg *config.DatabaseConfig usageCleanup config.UsageCleanupConfig encryptor SecretEncryptor - storeFactory BackupObjectStoreFactory - dumper DBDumper + // false 表示当前密钥由进程启动时临时生成，不能用于持久化可恢复的密文。 + encryptionKeyConfigured bool + storeFactory BackupObjectStoreFactory + dumper DBDumper opMu sync.Mutex // 保护 backingUp/restoring 标志 backingUp bool @@ -145,11 +157,14 @@ type BackupService struct { store BackupObjectStore s3Cfg *BackupS3Config - recordsMu sync.Mutex // 保护 records 的 load/save 操作 + recordsMu sync.Mutex // 保护 records 的 load/save 操作 + lifecycleMu sync.Mutex // 串行化新操作注册与 Stop/wg.Wait cronMu sync.Mutex cronSched *cron.Cron cronEntryID cron.EntryID + // cronCleanupEntryID 独立于 PostgreSQL 全库备份 schedule，始终按固定周期运行。 + cronCleanupEntryID cron.EntryID wg sync.WaitGroup // 追踪活跃的备份/恢复 goroutine shuttingDown atomic.Bool // 阻止新备份启动 @@ -167,19 +182,22 @@ func NewBackupService( bgCtx, bgCancel := context.WithCancel(context.Background()) dbCfg := &config.DatabaseConfig{} var usageCleanup config.UsageCleanupConfig + var encryptionKeyConfigured bool if cfg != nil { dbCfg = &cfg.Database usageCleanup = cfg.UsageCleanup + encryptionKeyConfigured = cfg.Totp.EncryptionKeyConfigured } return &BackupService{ - settingRepo: settingRepo, - dbCfg: dbCfg, - usageCleanup: usageCleanup, - encryptor: encryptor, - storeFactory: storeFactory, - dumper: dumper, - bgCtx: bgCtx, - bgCancel: bgCancel, + settingRepo: settingRepo, + dbCfg: dbCfg, + usageCleanup: usageCleanup, + encryptor: encryptor, + encryptionKeyConfigured: encryptionKeyConfigured, + storeFactory: storeFactory, + dumper: dumper, + bgCtx: bgCtx, + bgCancel: bgCancel, } } @@ -187,6 +205,11 @@ func NewBackupService( func (s *BackupService) Start() { s.cronSched = cron.New() s.cronSched.Start() + if err := s.applyExpirationCleanupSchedule(); err != nil { + logger.LegacyPrintf("service.backup", "[Backup] 注册归档过期清理任务失败: %v", err) + } else { + logger.LegacyPrintf("service.backup", "[Backup] 归档过期清理任务已启用: %s", backupExpirationCleanupCronExpr) + } // 清理重启后孤立的 running 记录 s.recoverStaleRecords() @@ -213,6 +236,7 @@ func (s *BackupService) recoverStaleRecords() { records, err := s.loadRecords(ctx) if err != nil { + logger.LegacyPrintf("service.backup", "[Backup] 加载孤立备份记录失败: %v", err) return } for i := range records { @@ -221,29 +245,42 @@ func (s *BackupService) recoverStaleRecords() { records[i].ErrorMsg = "interrupted by server restart" records[i].Progress = "" records[i].FinishedAt = time.Now().Format(time.RFC3339) - _ = s.saveRecord(ctx, &records[i]) - logger.LegacyPrintf("service.backup", "[Backup] recovered stale running record: %s", records[i].ID) + if err := s.saveRecord(ctx, &records[i]); err != nil { + logger.LegacyPrintf("service.backup", "[Backup] 标记孤立备份记录失败: id=%s err=%v", records[i].ID, err) + } else { + logger.LegacyPrintf("service.backup", "[Backup] recovered stale running record: %s", records[i].ID) + } } if records[i].RestoreStatus == "running" { records[i].RestoreStatus = "failed" records[i].RestoreError = "interrupted by server restart" - _ = s.saveRecord(ctx, &records[i]) - logger.LegacyPrintf("service.backup", "[Backup] recovered stale restoring record: %s", records[i].ID) + if err := s.saveRecord(ctx, &records[i]); err != nil { + logger.LegacyPrintf("service.backup", "[Backup] 标记孤立恢复记录失败: id=%s err=%v", records[i].ID, err) + } else { + logger.LegacyPrintf("service.backup", "[Backup] recovered stale restoring record: %s", records[i].ID) + } } } } // Stop 停止定时备份并等待活跃操作完成 func (s *BackupService) Stop() { + s.lifecycleMu.Lock() s.shuttingDown.Store(true) + s.lifecycleMu.Unlock() s.cronMu.Lock() if s.cronSched != nil { s.cronSched.Stop() } + s.cronEntryID = 0 + s.cronCleanupEntryID = 0 s.cronMu.Unlock() + if s.bgCancel != nil { + s.bgCancel() + } - // 等待活跃备份/恢复完成（最多 5 分钟） + // 后台 context 已取消；给流式上传/恢复一个有界退出窗口。 done := make(chan struct{}) go func() { s.wg.Wait() @@ -252,23 +289,18 @@ func (s *BackupService) Stop() { select { case <-done: logger.LegacyPrintf("service.backup", "[Backup] all active operations finished") - case <-time.After(5 * time.Minute): - logger.LegacyPrintf("service.backup", "[Backup] shutdown timeout after 5min, cancelling active operations") - if s.bgCancel != nil { - s.bgCancel() // 取消所有后台操作 - } - // 给 goroutine 时间响应取消并完成清理 - select { - case <-done: - logger.LegacyPrintf("service.backup", "[Backup] active operations cancelled and cleaned up") - case <-time.After(10 * time.Second): - logger.LegacyPrintf("service.backup", "[Backup] goroutine cleanup timed out") - } + case <-time.After(30 * time.Second): + logger.LegacyPrintf("service.backup", "[Backup] active operation shutdown timed out after 30s") } } // ─── S3 配置管理 ─── +// EncryptionKeyConfigured reports whether durable secrets can survive a restart. +func (s *BackupService) EncryptionKeyConfigured() bool { + return s != nil && s.encryptionKeyConfigured +} + func (s *BackupService) GetS3Config(ctx context.Context) (*BackupS3Config, error) { cfg, err := s.loadS3Config(ctx) if err != nil { @@ -283,19 +315,54 @@ func (s *BackupService) GetS3Config(ctx context.Context) (*BackupS3Config, error } func (s *BackupService) UpdateS3Config(ctx context.Context, cfg BackupS3Config) (*BackupS3Config, error) { - // 如果没提供 secret，保留原有值 + if !s.tryBeginRun() { + return nil, infraerrors.ServiceUnavailable("SERVER_SHUTTING_DOWN", "server is shutting down") + } + defer s.endRun() + if err := s.beginBackupCleanup(); err != nil { + return nil, err + } + defer s.finishBackupCleanup() + + old, err := s.loadS3Config(ctx) + if err != nil { + return nil, err + } + locationChanged := old == nil && backupStorageLocationConfigured(cfg) + if old != nil { + locationChanged = backupStorageLocationChanged(*old, cfg) + } + if locationChanged { + records, err := s.loadRecords(ctx) + if err != nil { + return nil, fmt.Errorf("load backup records before changing storage location: %w", err) + } + for _, record := range records { + if strings.TrimSpace(record.S3Key) != "" { + return nil, infraerrors.Conflict( + "BACKUP_STORAGE_LOCATION_IN_USE", + "cannot change backup endpoint, region, bucket, prefix, or path style while backup records still reference the current storage", + ) + } + } + } + + // 如果没提供 secret，保留原有明文值；落库前始终重新加密，避免只修改 + // 其他字段时把 loadS3Config 解密后的密钥以明文写回。 if cfg.SecretAccessKey == "" { - old, _ := s.loadS3Config(ctx) if old != nil { cfg.SecretAccessKey = old.SecretAccessKey } - } else { - // 加密 SecretAccessKey + } + if cfg.SecretAccessKey != "" { + if !s.encryptionKeyConfigured { + return nil, ErrSecretEncryptionKeyNotConfigured + } encrypted, err := s.encryptor.Encrypt(cfg.SecretAccessKey) if err != nil { return nil, fmt.Errorf("encrypt secret: %w", err) } - cfg.SecretAccessKey = encrypted + cfg.SecretAccessKey = backupEncryptedSecretV1 + encrypted } data, err := json.Marshal(cfg) @@ -316,10 +383,39 @@ func (s *BackupService) UpdateS3Config(ctx context.Context, cfg BackupS3Config) return &cfg, nil } +func backupStorageLocationChanged(oldCfg, newCfg BackupS3Config) bool { + normalizeEndpoint := func(value string) string { + return strings.TrimRight(strings.TrimSpace(value), "/") + } + normalizePrefix := func(value string) string { + value = strings.TrimRight(value, "/") + if value == "" { + return "backups" + } + return value + } + return normalizeEndpoint(oldCfg.Endpoint) != normalizeEndpoint(newCfg.Endpoint) || + strings.TrimSpace(oldCfg.Region) != strings.TrimSpace(newCfg.Region) || + strings.TrimSpace(oldCfg.Bucket) != strings.TrimSpace(newCfg.Bucket) || + normalizePrefix(oldCfg.Prefix) != normalizePrefix(newCfg.Prefix) || + oldCfg.ForcePathStyle != newCfg.ForcePathStyle +} + +func backupStorageLocationConfigured(cfg BackupS3Config) bool { + return strings.TrimSpace(cfg.Endpoint) != "" || + strings.TrimSpace(cfg.Region) != "" || + strings.TrimSpace(cfg.Bucket) != "" || + strings.TrimSpace(cfg.Prefix) != "" || + cfg.ForcePathStyle +} + func (s *BackupService) TestS3Connection(ctx context.Context, cfg BackupS3Config) error { // 如果没提供 secret，用已保存的 if cfg.SecretAccessKey == "" { - old, _ := s.loadS3Config(ctx) + old, err := s.loadS3Config(ctx) + if err != nil { + return err + } if old != nil { cfg.SecretAccessKey = old.SecretAccessKey } @@ -507,8 +603,10 @@ func (s *BackupService) removeCronSchedule() { } func (s *BackupService) runScheduledBackup() { - s.wg.Add(1) - defer s.wg.Done() + if !s.tryBeginRun() { + return + } + defer s.endRun() ctx, cancel := context.WithTimeout(s.bgCtx, 30*time.Minute) defer cancel() @@ -532,7 +630,7 @@ func (s *BackupService) runScheduledBackup() { } logger.LegacyPrintf("service.backup", "[Backup] 定时备份完成: id=%s size=%d", record.ID, record.SizeBytes) - // 清理过期备份（复用已加载的 schedule） + // 定时备份的份数/天数策略只适用于 PostgreSQL 全库备份。 if schedule == nil { return } @@ -541,20 +639,117 @@ func (s *BackupService) runScheduledBackup() { } } +func (s *BackupService) applyExpirationCleanupSchedule() error { + s.cronMu.Lock() + defer s.cronMu.Unlock() + + if s.cronSched == nil { + return fmt.Errorf("cron scheduler not initialized") + } + if s.cronCleanupEntryID != 0 { + s.cronSched.Remove(s.cronCleanupEntryID) + s.cronCleanupEntryID = 0 + } + entryID, err := s.cronSched.AddFunc(backupExpirationCleanupCronExpr, s.runScheduledExpirationCleanup) + if err != nil { + return fmt.Errorf("schedule backup expiration cleanup: %w", err) + } + s.cronCleanupEntryID = entryID + return nil +} + +func (s *BackupService) runScheduledExpirationCleanup() { + if !s.tryBeginRun() { + return + } + defer s.endRun() + + ctx, cancel := context.WithTimeout(s.bgCtx, 30*time.Minute) + defer cancel() + if err := s.cleanupExpiredBackupsWithRetry(ctx, backupExpirationCleanupAttempts, backupExpirationCleanupRetryWait); err != nil { + switch { + case errors.Is(err, ErrBackupInProgress): + logger.LegacyPrintf("service.backup", "[Backup] 归档过期清理跳过: 备份或归档正在运行") + case errors.Is(err, ErrRestoreInProgress): + logger.LegacyPrintf("service.backup", "[Backup] 归档过期清理跳过: 恢复正在运行") + default: + logger.LegacyPrintf("service.backup", "[Backup] 归档过期清理失败: %v", err) + } + return + } + logger.LegacyPrintf("service.backup", "[Backup] 归档过期清理执行完成") +} + +func (s *BackupService) cleanupExpiredBackupsWithRetry(ctx context.Context, attempts int, retryWait time.Duration) error { + if attempts <= 0 { + return fmt.Errorf("expiration cleanup attempts must be positive") + } + for attempt := 1; attempt <= attempts; attempt++ { + err := s.cleanupExpiredBackups(ctx) + if err == nil || (!errors.Is(err, ErrBackupInProgress) && !errors.Is(err, ErrRestoreInProgress)) { + return err + } + if attempt == attempts { + return fmt.Errorf("expiration cleanup remained busy after %d attempts: %w", attempts, err) + } + logger.LegacyPrintf( + "service.backup", + "[Backup] 归档过期清理等待重试: attempt=%d/%d wait=%s err=%v", + attempt, + attempts, + retryWait, + err, + ) + timer := time.NewTimer(retryWait) + select { + case <-ctx.Done(): + if !timer.Stop() { + select { + case <-timer.C: + default: + } + } + return ctx.Err() + case <-timer.C: + } + } + return nil +} + +// tryBeginRun 将所有同步、异步和定时操作的 wg.Add 与 Stop/wg.Wait 串行化。 +func (s *BackupService) tryBeginRun() bool { + s.lifecycleMu.Lock() + defer s.lifecycleMu.Unlock() + if s.shuttingDown.Load() { + return false + } + s.wg.Add(1) + return true +} + +func (s *BackupService) endRun() { + s.wg.Done() +} + // ─── 备份/恢复核心 ─── // CreateBackup 创建全量数据库备份并上传到 S3（流式处理） // expireDays: 备份过期天数，0=永不过期，默认14天 func (s *BackupService) CreateBackup(ctx context.Context, triggeredBy string, expireDays int) (*BackupRecord, error) { - if s.shuttingDown.Load() { + if !s.tryBeginRun() { return nil, infraerrors.ServiceUnavailable("SERVER_SHUTTING_DOWN", "server is shutting down") } + defer s.endRun() s.opMu.Lock() if s.backingUp { s.opMu.Unlock() return nil, ErrBackupInProgress } + if s.restoring { + s.opMu.Unlock() + return nil, ErrRestoreInProgress + } s.backingUp = true s.opMu.Unlock() defer func() { @@ -603,7 +798,9 @@ func (s *BackupService) CreateBackup(ctx context.Context, triggeredBy string, ex record.Status = "failed" record.ErrorMsg = fmt.Sprintf("pg_dump failed: %v", err) record.FinishedAt = time.Now().Format(time.RFC3339) - _ = s.saveRecord(ctx, record) + if saveErr := s.saveRecord(ctx, record); saveErr != nil { + logger.LegacyPrintf("service.backup", "[Backup] 保存失败备份记录失败: %v", saveErr) + } return record, fmt.Errorf("pg_dump: %w", err) } @@ -646,16 +843,24 @@ func (s *BackupService) CreateBackup(ctx context.Context, triggeredBy string, ex } record.ErrorMsg = errMsg record.FinishedAt = time.Now().Format(time.RFC3339) - _ = s.saveRecord(ctx, record) + if saveErr := s.saveRecord(ctx, record); saveErr != nil { + logger.LegacyPrintf("service.backup", "[Backup] 保存上传失败备份记录失败: %v", saveErr) + } return record, fmt.Errorf("backup upload: %w", err) } - <-gzipDone // 确保 gzip goroutine 已退出 + if gzErr := <-gzipDone; gzErr != nil { + record.Status = "failed" + record.ErrorMsg = fmt.Sprintf("gzip/dump failed: %v", gzErr) + record.FinishedAt = time.Now().Format(time.RFC3339) + archiveErr := fmt.Errorf("backup gzip/dump: %w", gzErr) + return record, compensateUploadedObject(ctx, objectStore, record.S3Key, archiveErr) + } record.SizeBytes = sizeBytes record.Status = "completed" record.FinishedAt = time.Now().Format(time.RFC3339) - if err := s.saveRecord(ctx, record); err != nil { - logger.LegacyPrintf("service.backup", "[Backup] 保存备份记录失败: %v", err) + if err := s.saveCompletedRecordOrCompensate(ctx, objectStore, record); err != nil { + return record, err } return record, nil @@ -701,15 +906,20 @@ func (s *BackupService) CreateDataArchive(ctx context.Context, input DataArchive if input.TriggeredBy == "" { input.TriggeredBy = "system" } - if s.shuttingDown.Load() { + if !s.tryBeginRun() { return nil, infraerrors.ServiceUnavailable("SERVER_SHUTTING_DOWN", "server is shutting down") } + defer s.endRun() s.opMu.Lock() if s.backingUp { s.opMu.Unlock() return nil, ErrBackupInProgress } + if s.restoring { + s.opMu.Unlock() + return nil, ErrRestoreInProgress + } s.backingUp = true s.opMu.Unlock() defer func() { @@ -778,35 +988,47 @@ func (s *BackupService) CreateDataArchive(ctx context.Context, input DataArchive record.ErrorMsg = fmt.Sprintf("gzip/archive failed: %v", gzErr) } record.FinishedAt = time.Now().Format(time.RFC3339) - _ = s.saveRecord(ctx, record) + if saveErr := s.saveRecord(ctx, record); saveErr != nil { + logger.LegacyPrintf("service.backup", "[Backup] 保存归档上传失败状态失败: %v", saveErr) + } return record, fmt.Errorf("%s archive upload: %w", input.BackupType, err) } if gzErr := <-gzipDone; gzErr != nil { record.Status = "failed" record.ErrorMsg = fmt.Sprintf("gzip/archive failed: %v", gzErr) record.FinishedAt = time.Now().Format(time.RFC3339) - _ = s.saveRecord(ctx, record) - return record, gzErr + archiveErr := fmt.Errorf("%s archive gzip: %w", input.BackupType, gzErr) + return record, compensateUploadedObject(ctx, objectStore, record.S3Key, archiveErr) } record.SizeBytes = sizeBytes record.Status = "completed" record.FinishedAt = time.Now().Format(time.RFC3339) - if err := s.saveRecord(ctx, record); err != nil { - logger.LegacyPrintf("service.backup", "[Backup] failed to save %s archive record: %v", input.BackupType, err) + if err := s.saveCompletedRecordOrCompensate(ctx, objectStore, record); err != nil { + return record, err } return record, nil } func (s *BackupService) StartBackup(ctx context.Context, triggeredBy string, expireDays int) (*BackupRecord, error) { - if s.shuttingDown.Load() { + if !s.tryBeginRun() { return nil, infraerrors.ServiceUnavailable("SERVER_SHUTTING_DOWN", "server is shutting down") } + runOwned := true + defer func() { + if runOwned { + s.endRun() + } + }() s.opMu.Lock() if s.backingUp { s.opMu.Unlock() return nil, ErrBackupInProgress } + if s.restoring { + s.opMu.Unlock() + return nil, ErrRestoreInProgress + } s.backingUp = true s.opMu.Unlock() @@ -864,9 +1086,8 @@ func (s *BackupService) StartBackup(ctx context.Context, triggeredBy string, exp // 在启动 goroutine 前完成拷贝，避免数据竞争 result := *record - s.wg.Add(1) go func() { - defer s.wg.Done() + defer s.endRun() defer func() { s.opMu.Lock() s.backingUp = false @@ -879,11 +1100,14 @@ func (s *BackupService) StartBackup(ctx context.Context, triggeredBy string, exp record.ErrorMsg = fmt.Sprintf("internal panic: %v", r) record.Progress = "" record.FinishedAt = time.Now().Format(time.RFC3339) - _ = s.saveRecord(context.Background(), record) + if saveErr := s.saveRecord(context.Background(), record); saveErr != nil { + logger.LegacyPrintf("service.backup", "[Backup] 保存 panic 失败状态失败: %v", saveErr) + } } }() s.executeBackup(record, objectStore) }() + runOwned = false return &result, nil } @@ -895,7 +1119,9 @@ func (s *BackupService) executeBackup(record *BackupRecord, objectStore BackupOb // 阶段1: pg_dump record.Progress = "dumping" - _ = s.saveRecord(ctx, record) + if saveErr := s.saveRecord(ctx, record); saveErr != nil { + logger.LegacyPrintf("service.backup", "[Backup] 保存备份进度 dumping 失败: %v", saveErr) + } dumpReader, err := s.dumper.Dump(ctx) if err != nil { @@ -903,13 +1129,17 @@ func (s *BackupService) executeBackup(record *BackupRecord, objectStore BackupOb record.ErrorMsg = fmt.Sprintf("pg_dump failed: %v", err) record.Progress = "" record.FinishedAt = time.Now().Format(time.RFC3339) - _ = s.saveRecord(context.Background(), record) + if saveErr := s.saveRecord(context.Background(), record); saveErr != nil { + logger.LegacyPrintf("service.backup", "[Backup] 保存异步 pg_dump 失败状态失败: %v", saveErr) + } return } // 阶段2: gzip + upload record.Progress = "uploading" - _ = s.saveRecord(ctx, record) + if saveErr := s.saveRecord(ctx, record); saveErr != nil { + logger.LegacyPrintf("service.backup", "[Backup] 保存备份进度 uploading 失败: %v", saveErr) + } pr, pw := io.Pipe() gzipDone := make(chan error, 1) @@ -950,23 +1180,49 @@ func (s *BackupService) executeBackup(record *BackupRecord, objectStore BackupOb record.ErrorMsg = errMsg record.Progress = "" record.FinishedAt = time.Now().Format(time.RFC3339) - _ = s.saveRecord(context.Background(), record) + if saveErr := s.saveRecord(context.Background(), record); saveErr != nil { + logger.LegacyPrintf("service.backup", "[Backup] 保存异步上传失败状态失败: %v", saveErr) + } + return + } + if gzErr := <-gzipDone; gzErr != nil { + record.Status = "failed" + record.ErrorMsg = fmt.Sprintf("gzip/dump failed: %v", gzErr) + record.Progress = "" + record.FinishedAt = time.Now().Format(time.RFC3339) + cleanupErr := compensateUploadedObject(ctx, objectStore, record.S3Key, fmt.Errorf("backup gzip/dump: %w", gzErr)) + record.ErrorMsg = cleanupErr.Error() + if saveErr := s.saveRecord(context.Background(), record); saveErr != nil { + logger.LegacyPrintf("service.backup", "[Backup] 保存异步 gzip 失败状态失败: %v", saveErr) + } return } - <-gzipDone // 确保 gzip goroutine 已退出 record.SizeBytes = sizeBytes record.Status = "completed" record.Progress = "" record.FinishedAt = time.Now().Format(time.RFC3339) - if err := s.saveRecord(context.Background(), record); err != nil { - logger.LegacyPrintf("service.backup", "[Backup] 保存备份记录失败: %v", err) + if err := s.saveCompletedRecordOrCompensate(context.Background(), objectStore, record); err != nil { + logger.LegacyPrintf("service.backup", "[Backup] 保存备份记录失败，已将备份标记为失败并尝试补偿: %v", err) + // 初始 running 记录已经存在；若存储故障短暂恢复，尽力持久化失败状态。 + if saveErr := s.saveRecord(context.Background(), record); saveErr != nil { + logger.LegacyPrintf("service.backup", "[Backup] 保存补偿失败状态失败: %v", saveErr) + } } } // RestoreBackup 从 S3 下载备份并流式恢复到数据库 func (s *BackupService) RestoreBackup(ctx context.Context, backupID string) error { + if !s.tryBeginRun() { + return infraerrors.ServiceUnavailable("SERVER_SHUTTING_DOWN", "server is shutting down") + } + defer s.endRun() + s.opMu.Lock() + if s.backingUp { + s.opMu.Unlock() + return ErrBackupInProgress + } if s.restoring { s.opMu.Unlock() return ErrRestoreInProgress @@ -1023,11 +1279,21 @@ func (s *BackupService) RestoreBackup(ctx context.Context, backupID string) erro // StartRestore 异步恢复备份，立即返回 func (s *BackupService) StartRestore(ctx context.Context, backupID string) (*BackupRecord, error) { - if s.shuttingDown.Load() { + if !s.tryBeginRun() { return nil, infraerrors.ServiceUnavailable("SERVER_SHUTTING_DOWN", "server is shutting down") } + runOwned := true + defer func() { + if runOwned { + s.endRun() + } + }() s.opMu.Lock() + if s.backingUp { + s.opMu.Unlock() + return nil, ErrBackupInProgress + } if s.restoring { s.opMu.Unlock() return nil, ErrRestoreInProgress @@ -1066,14 +1332,15 @@ func (s *BackupService) StartRestore(ctx context.Context, backupID string) (*Bac } record.RestoreStatus = "running" - _ = s.saveRecord(ctx, record) + if err := s.saveRecord(ctx, record); err != nil { + return nil, fmt.Errorf("save restore status: %w", err) + } launched = true result := *record - s.wg.Add(1) go func() { - defer s.wg.Done() + defer s.endRun() defer func() { s.opMu.Lock() s.restoring = false @@ -1084,11 +1351,14 @@ func (s *BackupService) StartRestore(ctx context.Context, backupID string) (*Bac logger.LegacyPrintf("service.backup", "[Backup] restore panic recovered: %v", r) record.RestoreStatus = "failed" record.RestoreError = fmt.Sprintf("internal panic: %v", r) - _ = s.saveRecord(context.Background(), record) + if saveErr := s.saveRecord(context.Background(), record); saveErr != nil { + logger.LegacyPrintf("service.backup", "[Backup] 保存恢复 panic 失败状态失败: %v", saveErr) + } } }() s.executeRestore(record, objectStore) }() + runOwned = false return &result, nil } @@ -1102,7 +1372,9 @@ func (s *BackupService) executeRestore(record *BackupRecord, objectStore BackupO if err != nil { record.RestoreStatus = "failed" record.RestoreError = fmt.Sprintf("S3 download failed: %v", err) - _ = s.saveRecord(context.Background(), record) + if saveErr := s.saveRecord(context.Background(), record); saveErr != nil { + logger.LegacyPrintf("service.backup", "[Backup] 保存恢复下载失败状态失败: %v", saveErr) + } return } defer func() { _ = body.Close() }() @@ -1111,7 +1383,9 @@ func (s *BackupService) executeRestore(record *BackupRecord, objectStore BackupO if err != nil { record.RestoreStatus = "failed" record.RestoreError = fmt.Sprintf("gzip reader: %v", err) - _ = s.saveRecord(context.Background(), record) + if saveErr := s.saveRecord(context.Background(), record); saveErr != nil { + logger.LegacyPrintf("service.backup", "[Backup] 保存恢复 gzip 失败状态失败: %v", saveErr) + } return } defer func() { _ = gzReader.Close() }() @@ -1119,7 +1393,9 @@ func (s *BackupService) executeRestore(record *BackupRecord, objectStore BackupO if err := s.dumper.Restore(ctx, gzReader); err != nil { record.RestoreStatus = "failed" record.RestoreError = fmt.Sprintf("pg restore: %v", err) - _ = s.saveRecord(context.Background(), record) + if saveErr := s.saveRecord(context.Background(), record); saveErr != nil { + logger.LegacyPrintf("service.backup", "[Backup] 保存恢复失败状态失败: %v", saveErr) + } return } @@ -1158,6 +1434,16 @@ func (s *BackupService) GetBackupRecord(ctx context.Context, backupID string) (* } func (s *BackupService) DeleteBackup(ctx context.Context, backupID string) error { + if !s.tryBeginRun() { + return infraerrors.ServiceUnavailable("SERVER_SHUTTING_DOWN", "server is shutting down") + } + defer s.endRun() + + if err := s.beginBackupCleanup(); err != nil { + return err + } + defer s.finishBackupCleanup() + s.recordsMu.Lock() defer s.recordsMu.Unlock() @@ -1179,14 +1465,10 @@ func (s *BackupService) DeleteBackup(ctx context.Context, backupID string) error return ErrBackupNotFound } - // 从 S3 删除 - if found.S3Key != "" && found.Status == "completed" { - s3Cfg, err := s.loadS3Config(ctx) - if err == nil && s3Cfg != nil && s3Cfg.IsConfigured() { - objectStore, err := s.getOrCreateStore(ctx, s3Cfg) - if err == nil { - _ = objectStore.Delete(ctx, found.S3Key) - } + // 有对象键时必须先确认 S3 删除成功，防止对象仍存在但记录已经丢失。 + if strings.TrimSpace(found.S3Key) != "" { + if err := s.deleteS3Object(ctx, found.S3Key); err != nil { + return fmt.Errorf("delete backup object for record %s: %w", found.ID, err) } } @@ -1223,21 +1505,35 @@ func (s *BackupService) GetBackupDownloadURL(ctx context.Context, backupID strin func (s *BackupService) loadS3Config(ctx context.Context) (*BackupS3Config, error) { raw, err := s.settingRepo.GetValue(ctx, settingKeyBackupS3Config) - if err != nil || raw == "" { + if errors.Is(err, ErrSettingNotFound) { return nil, nil //nolint:nilnil // no config is a valid state } + if err != nil { + return nil, fmt.Errorf("load S3 config setting: %w", err) + } + if strings.TrimSpace(raw) == "" { + return nil, ErrBackupS3ConfigCorrupt + } var cfg BackupS3Config if err := json.Unmarshal([]byte(raw), &cfg); err != nil { return nil, ErrBackupS3ConfigCorrupt } - // 解密 SecretAccessKey + // 新格式带显式版本标记：解密失败说明密钥变化或数据损坏，必须快速失败， + // 不能把密文继续当成 SecretAccessKey 传给对象存储。 if cfg.SecretAccessKey != "" { - decrypted, err := s.encryptor.Decrypt(cfg.SecretAccessKey) - if err != nil { - // 兼容未加密的旧数据：如果解密失败，保持原值 - logger.LegacyPrintf("service.backup", "[Backup] S3 SecretAccessKey 解密失败（可能是旧的未加密数据）: %v", err) - } else { + if strings.HasPrefix(cfg.SecretAccessKey, backupEncryptedSecretV1) { + ciphertext := strings.TrimPrefix(cfg.SecretAccessKey, backupEncryptedSecretV1) + decrypted, err := s.encryptor.Decrypt(ciphertext) + if err != nil { + return nil, fmt.Errorf("%w: decrypt S3 secret access key: %v", ErrBackupS3ConfigCorrupt, err) + } cfg.SecretAccessKey = decrypted + } else if decrypted, err := s.encryptor.Decrypt(cfg.SecretAccessKey); err == nil { + // 兼容旧版无标记密文；下一次保存配置时会迁移到 v1 标记格式。 + cfg.SecretAccessKey = decrypted + } else { + // 兼容早期直接保存的明文。只有无版本标记的数据允许走此分支。 + logger.LegacyPrintf("service.backup", "[Backup] 检测到旧版未标记的 S3 SecretAccessKey，将在下次保存时迁移") } } return &cfg, nil @@ -1282,9 +1578,15 @@ func (s *BackupService) loadRecords(ctx context.Context) ([]BackupRecord, error) // loadRecordsLocked 在已持有 recordsMu 锁的情况下加载记录 func (s *BackupService) loadRecordsLocked(ctx context.Context) ([]BackupRecord, error) { raw, err := s.settingRepo.GetValue(ctx, settingKeyBackupRecords) - if err != nil || raw == "" { + if errors.Is(err, ErrSettingNotFound) { return nil, nil //nolint:nilnil // no records is a valid state } + if err != nil { + return nil, fmt.Errorf("load backup records setting: %w", err) + } + if strings.TrimSpace(raw) == "" { + return nil, ErrBackupRecordsCorrupt + } var records []BackupRecord if err := json.Unmarshal([]byte(raw), &records); err != nil { return nil, ErrBackupRecordsCorrupt @@ -1306,7 +1608,10 @@ func (s *BackupService) saveRecord(ctx context.Context, record *BackupRecord) er s.recordsMu.Lock() defer s.recordsMu.Unlock() - records, _ := s.loadRecordsLocked(ctx) + records, err := s.loadRecordsLocked(ctx) + if err != nil { + return err + } // 更新已有记录或追加 found := false @@ -1320,19 +1625,124 @@ func (s *BackupService) saveRecord(ctx context.Context, record *BackupRecord) er if !found { records = append(records, *record) } + return s.saveRecordsLocked(ctx, records) +} - // 限制记录数量 - if len(records) > maxBackupRecords { - records = records[len(records)-maxBackupRecords:] +func (s *BackupService) saveCompletedRecordOrCompensate( + ctx context.Context, + objectStore BackupObjectStore, + record *BackupRecord, +) error { + if err := s.saveRecord(ctx, record); err != nil { + saveErr := fmt.Errorf("save completed %s backup record: %w", record.BackupType, err) + record.Status = "failed" + record.ErrorMsg = saveErr.Error() + return compensateUploadedObject(ctx, objectStore, record.S3Key, saveErr) } + return nil +} - return s.saveRecordsLocked(ctx, records) +func compensateUploadedObject( + ctx context.Context, + objectStore BackupObjectStore, + key string, + cause error, +) error { + compensationCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), backupCompensationTimeout) + defer cancel() + if err := objectStore.Delete(compensationCtx, key); err != nil { + return errors.Join(cause, fmt.Errorf("delete uploaded object %s after failure: %w", key, err)) + } + return cause +} + +func (s *BackupService) cleanupExpiredBackups(ctx context.Context) error { + if err := s.beginBackupCleanup(); err != nil { + return err + } + defer s.finishBackupCleanup() + + return s.cleanupExpiredBackupRecords(ctx, time.Now()) +} + +// cleanupExpiredBackupRecords 仅按 ExpiresAt 清理已授权的数据归档类型。 +// 调用方必须已经占用备份操作槽位，避免与备份、归档及恢复并发。 +func (s *BackupService) cleanupExpiredBackupRecords(ctx context.Context, now time.Time) error { + s.recordsMu.Lock() + defer s.recordsMu.Unlock() + + records, err := s.loadRecordsLocked(ctx) + if err != nil { + return err + } + + toDelete := make(map[string]struct{}) + var cleanupErr error + for _, record := range records { + eligible, typeErr := isExpirationManagedArchive(record) + if typeErr != nil { + cleanupErr = errors.Join(cleanupErr, typeErr) + continue + } + if !eligible { + continue + } + expiresAtValue := strings.TrimSpace(record.ExpiresAt) + if expiresAtValue == "" { + continue + } + expiresAt, err := time.Parse(time.RFC3339, expiresAtValue) + if err != nil { + cleanupErr = errors.Join(cleanupErr, fmt.Errorf("parse expires_at for backup record %s: %w", record.ID, err)) + continue + } + if now.Before(expiresAt) { + continue + } + if record.Status == "running" || record.RestoreStatus == "running" { + cleanupErr = errors.Join(cleanupErr, fmt.Errorf("backup record %s is expired but still active", record.ID)) + continue + } + toDelete[record.ID] = struct{}{} + if len(toDelete) >= backupExpirationCleanupBatchSize { + break + } + } + + deleted, remaining, deleteErr := s.deleteBackupRecordObjectsLocked(ctx, records, toDelete) + cleanupErr = errors.Join(cleanupErr, deleteErr) + if deleted == 0 { + return cleanupErr + } + if err := s.saveRecordsLocked(ctx, remaining); err != nil { + return errors.Join(cleanupErr, fmt.Errorf("save records after expiration cleanup: %w", err)) + } + logger.LegacyPrintf("service.backup", "[Backup] 自动清理了 %d 个已过期归档", deleted) + return cleanupErr +} + +func isExpirationManagedArchive(record BackupRecord) (bool, error) { + switch strings.TrimSpace(record.BackupType) { + case "usage_logs_archive", "ops_system_logs_archive", "ops_error_logs_archive": + return true, nil + case "", "postgres": + return false, nil + default: + if strings.TrimSpace(record.ExpiresAt) == "" { + return false, nil + } + return false, fmt.Errorf("backup record %s has unsupported expiration-managed type %q", record.ID, record.BackupType) + } } func (s *BackupService) cleanupOldBackups(ctx context.Context, schedule *BackupScheduleConfig) error { if schedule == nil { return nil } + if err := s.beginBackupCleanup(); err != nil { + return err + } + defer s.finishBackupCleanup() s.recordsMu.Lock() defer s.recordsMu.Unlock() @@ -1342,15 +1752,46 @@ func (s *BackupService) cleanupOldBackups(ctx context.Context, schedule *BackupS return err } - // 按时间倒序 - sort.Slice(records, func(i, j int) bool { - return records[i].StartedAt > records[j].StartedAt - }) + if schedule.RetainCount <= 0 && schedule.RetainDays <= 0 { + return nil + } - var toDelete []BackupRecord - var toKeep []BackupRecord + // 份数和保留天数只针对已完成的 PostgreSQL 全库备份计算，归档记录不参与排序。 + type postgresBackupWithTime struct { + record BackupRecord + startedAt time.Time + } + postgresRecords := make([]postgresBackupWithTime, 0, len(records)) + var cleanupErr error + for _, record := range records { + if !isPostgresBackupRecord(record) || record.Status != "completed" { + continue + } + startedAt, err := time.Parse(time.RFC3339, strings.TrimSpace(record.StartedAt)) + if err != nil { + cleanupErr = errors.Join(cleanupErr, fmt.Errorf("parse started_at for postgres backup record %s: %w", record.ID, err)) + continue + } + postgresRecords = append(postgresRecords, postgresBackupWithTime{record: record, startedAt: startedAt}) + } + // RetainCount 依赖完整且可靠的时间顺序；存在损坏记录时直接失败，禁止猜测排序后误删。 + if cleanupErr != nil { + return cleanupErr + } + sort.Slice(postgresRecords, func(i, j int) bool { + return postgresRecords[i].startedAt.After(postgresRecords[j].startedAt) + }) - for i, r := range records { + toDelete := make(map[string]struct{}) + cutoff := time.Time{} + if schedule.RetainDays > 0 { + cutoff = time.Now().AddDate(0, 0, -schedule.RetainDays) + } + for i, candidate := range postgresRecords { + // 无论份数/天数策略如何，始终保留最新一份成功的全库备份。 + if i == 0 { + continue + } shouldDelete := false // 按保留份数清理 @@ -1359,38 +1800,91 @@ func (s *BackupService) cleanupOldBackups(ctx context.Context, schedule *BackupS } // 按保留天数清理 - if schedule.RetainDays > 0 && r.StartedAt != "" { - startedAt, err := time.Parse(time.RFC3339, r.StartedAt) - if err == nil && time.Since(startedAt) > time.Duration(schedule.RetainDays)*24*time.Hour { - shouldDelete = true - } + if schedule.RetainDays > 0 && candidate.startedAt.Before(cutoff) { + shouldDelete = true } - if shouldDelete && r.Status == "completed" { - toDelete = append(toDelete, r) - } else { - toKeep = append(toKeep, r) + if shouldDelete { + toDelete[candidate.record.ID] = struct{}{} } } - // 删除 S3 上的文件 - for _, r := range toDelete { - if r.S3Key != "" { - _ = s.deleteS3Object(ctx, r.S3Key) - } + deleted, remaining, deleteErr := s.deleteBackupRecordObjectsLocked(ctx, records, toDelete) + cleanupErr = errors.Join(cleanupErr, deleteErr) + if deleted == 0 { + return cleanupErr + } + if err := s.saveRecordsLocked(ctx, remaining); err != nil { + return errors.Join(cleanupErr, fmt.Errorf("save records after postgres backup cleanup: %w", err)) } + logger.LegacyPrintf("service.backup", "[Backup] 自动清理了 %d 个 PostgreSQL 全库备份", deleted) + return cleanupErr +} - if len(toDelete) > 0 { - logger.LegacyPrintf("service.backup", "[Backup] 自动清理了 %d 个过期备份", len(toDelete)) - return s.saveRecordsLocked(ctx, toKeep) +func (s *BackupService) beginBackupCleanup() error { + if s.shuttingDown.Load() { + return infraerrors.ServiceUnavailable("SERVER_SHUTTING_DOWN", "server is shutting down") } + s.opMu.Lock() + defer s.opMu.Unlock() + if s.backingUp { + return ErrBackupInProgress + } + if s.restoring { + return ErrRestoreInProgress + } + s.backingUp = true return nil } +func (s *BackupService) finishBackupCleanup() { + s.opMu.Lock() + s.backingUp = false + s.opMu.Unlock() +} + +func isPostgresBackupRecord(record BackupRecord) bool { + backupType := strings.TrimSpace(record.BackupType) + return backupType == "" || backupType == "postgres" +} + +// deleteBackupRecordObjectsLocked 删除候选记录对应的对象。 +// 只有对象删除成功的记录才会从 remaining 中移除；失败项保留并聚合返回错误。 +func (s *BackupService) deleteBackupRecordObjectsLocked( + ctx context.Context, + records []BackupRecord, + toDelete map[string]struct{}, +) (int, []BackupRecord, error) { + remaining := make([]BackupRecord, 0, len(records)) + deleted := 0 + var cleanupErr error + for _, record := range records { + if _, ok := toDelete[record.ID]; !ok { + remaining = append(remaining, record) + continue + } + if strings.TrimSpace(record.S3Key) == "" { + remaining = append(remaining, record) + cleanupErr = errors.Join(cleanupErr, fmt.Errorf("backup record %s cannot be deleted: S3 key is empty", record.ID)) + continue + } + if err := s.deleteS3Object(ctx, record.S3Key); err != nil { + remaining = append(remaining, record) + cleanupErr = errors.Join(cleanupErr, fmt.Errorf("delete S3 object for backup record %s: %w", record.ID, err)) + continue + } + deleted++ + } + return deleted, remaining, cleanupErr +} + func (s *BackupService) deleteS3Object(ctx context.Context, key string) error { s3Cfg, err := s.loadS3Config(ctx) - if err != nil || s3Cfg == nil { - return nil + if err != nil { + return fmt.Errorf("load S3 config: %w", err) + } + if s3Cfg == nil || !s3Cfg.IsConfigured() { + return ErrBackupS3NotConfigured } objectStore, err := s.getOrCreateStore(ctx, s3Cfg) if err != nil { diff --git a/backend/internal/service/backup_service_test.go b/backend/internal/service/backup_service_test.go index 31c5ec15c..2ee51ae8b 100644 --- a/backend/internal/service/backup_service_test.go +++ b/backend/internal/service/backup_service_test.go @@ -23,12 +23,18 @@ import ( // ─── Mocks ─── type mockSettingRepo struct { - mu sync.Mutex - data map[string]string + mu sync.Mutex + data map[string]string + getValueErrors map[string]error + setErrors map[string]error } func newMockSettingRepo() *mockSettingRepo { - return &mockSettingRepo{data: make(map[string]string)} + return &mockSettingRepo{ + data: make(map[string]string), + getValueErrors: make(map[string]error), + setErrors: make(map[string]error), + } } func (m *mockSettingRepo) Get(_ context.Context, key string) (*Setting, error) { @@ -44,9 +50,12 @@ func (m *mockSettingRepo) Get(_ context.Context, key string) (*Setting, error) { func (m *mockSettingRepo) GetValue(_ context.Context, key string) (string, error) { m.mu.Lock() defer m.mu.Unlock() + if err := m.getValueErrors[key]; err != nil { + return "", err + } v, ok := m.data[key] if !ok { - return "", nil + return "", ErrSettingNotFound } return v, nil } @@ -54,6 +63,9 @@ func (m *mockSettingRepo) GetValue(_ context.Context, key string) (string, error func (m *mockSettingRepo) Set(_ context.Context, key, value string) error { m.mu.Lock() defer m.mu.Unlock() + if err := m.setErrors[key]; err != nil { + return err + } m.data[key] = value return nil } @@ -161,12 +173,22 @@ func (d *blockingDumper) Restore(_ context.Context, data io.Reader) error { } type mockObjectStore struct { - objects map[string][]byte - mu sync.Mutex + objects map[string][]byte + uploadHook func() + deleteErr error + deleteErrors map[string]error + deleteCalls []string + deleteBlock <-chan struct{} + deleteStarted chan struct{} + deleteStartedOnce sync.Once + mu sync.Mutex } func newMockObjectStore() *mockObjectStore { - return &mockObjectStore{objects: make(map[string][]byte)} + return &mockObjectStore{ + objects: make(map[string][]byte), + deleteErrors: make(map[string]error), + } } func (m *mockObjectStore) Upload(_ context.Context, key string, body io.Reader, _ string) (int64, error) { @@ -176,7 +198,11 @@ func (m *mockObjectStore) Upload(_ context.Context, key string, body io.Reader, } m.mu.Lock() m.objects[key] = data + hook := m.uploadHook m.mu.Unlock() + if hook != nil { + hook() + } return int64(len(data)), nil } @@ -190,7 +216,31 @@ func (m *mockObjectStore) Download(_ context.Context, key string) (io.ReadCloser return io.NopCloser(bytes.NewReader(data)), nil } -func (m *mockObjectStore) Delete(_ context.Context, key string) error { +func (m *mockObjectStore) Delete(ctx context.Context, key string) error { + m.mu.Lock() + m.deleteCalls = append(m.deleteCalls, key) + deleteErr := m.deleteErrors[key] + if deleteErr == nil { + deleteErr = m.deleteErr + } + deleteBlock := m.deleteBlock + deleteStarted := m.deleteStarted + m.mu.Unlock() + + if deleteStarted != nil { + m.deleteStartedOnce.Do(func() { close(deleteStarted) }) + } + if deleteBlock != nil { + select { + case <-deleteBlock: + case <-ctx.Done(): + return ctx.Err() + } + } + if deleteErr != nil { + return deleteErr + } + m.mu.Lock() delete(m.objects, key) m.mu.Unlock() @@ -213,6 +263,7 @@ func newTestBackupService(repo *mockSettingRepo, dumper DBDumper, store *mockObj User: "test", DBName: "testdb", }, + Totp: config.TotpConfig{EncryptionKeyConfigured: true}, } factory := func(_ context.Context, _ *BackupS3Config) (BackupObjectStore, error) { return store, nil @@ -220,6 +271,21 @@ func newTestBackupService(repo *mockSettingRepo, dumper DBDumper, store *mockObj return NewBackupService(repo, cfg, &plainEncryptor{}, factory, dumper) } +func newTestBackupServiceWithEphemeralKey(repo *mockSettingRepo) *BackupService { + cfg := &config.Config{ + Database: config.DatabaseConfig{ + Host: "localhost", + Port: 5432, + User: "test", + DBName: "testdb", + }, + Totp: config.TotpConfig{EncryptionKeyConfigured: false}, + } + return NewBackupService(repo, cfg, &plainEncryptor{}, func(_ context.Context, _ *BackupS3Config) (BackupObjectStore, error) { + return newMockObjectStore(), nil + }, &mockDumper{}) +} + func seedS3Config(t *testing.T, repo *mockSettingRepo) { t.Helper() cfg := BackupS3Config{ @@ -232,6 +298,17 @@ func seedS3Config(t *testing.T, repo *mockSettingRepo) { require.NoError(t, repo.Set(context.Background(), settingKeyBackupS3Config, string(data))) } +func seedBackupRecord(t *testing.T, svc *BackupService, store *mockObjectStore, record BackupRecord) { + t.Helper() + require.NoError(t, svc.saveRecord(context.Background(), &record)) + if record.S3Key == "" { + return + } + store.mu.Lock() + store.objects[record.S3Key] = []byte(record.ID) + store.mu.Unlock() +} + // ─── Tests ─── func TestBackupService_S3ConfigEncryption(t *testing.T) { @@ -251,7 +328,7 @@ func TestBackupService_S3ConfigEncryption(t *testing.T) { raw, _ := repo.GetValue(context.Background(), settingKeyBackupS3Config) var stored BackupS3Config require.NoError(t, json.Unmarshal([]byte(raw), &stored)) - require.Equal(t, "ENC:my-secret", stored.SecretAccessKey) + require.Equal(t, backupEncryptedSecretV1+"ENC:my-secret", stored.SecretAccessKey) // 通过 GetS3Config 获取应该脱敏 cfg, err := svc.GetS3Config(context.Background()) @@ -265,6 +342,18 @@ func TestBackupService_S3ConfigEncryption(t *testing.T) { require.Equal(t, "my-secret", internal.SecretAccessKey) } +func TestBackupService_S3ConfigRejectsEphemeralEncryptionKey(t *testing.T) { + svc := newTestBackupServiceWithEphemeralKey(newMockSettingRepo()) + + _, err := svc.UpdateS3Config(context.Background(), BackupS3Config{ + Bucket: "my-bucket", + AccessKeyID: "AKID", + SecretAccessKey: "my-secret", + }) + require.ErrorIs(t, err, ErrSecretEncryptionKeyNotConfigured) + require.False(t, svc.EncryptionKeyConfigured()) +} + func TestBackupService_S3ConfigKeepExistingSecret(t *testing.T) { repo := newMockSettingRepo() svc := newTestBackupService(repo, &mockDumper{}, newMockObjectStore()) @@ -288,6 +377,128 @@ func TestBackupService_S3ConfigKeepExistingSecret(t *testing.T) { require.NoError(t, err) require.Equal(t, "original-secret", internal.SecretAccessKey) require.Equal(t, "AKID-NEW", internal.AccessKeyID) + + raw, err := repo.GetValue(context.Background(), settingKeyBackupS3Config) + require.NoError(t, err) + var stored BackupS3Config + require.NoError(t, json.Unmarshal([]byte(raw), &stored)) + require.Equal(t, backupEncryptedSecretV1+"ENC:original-secret", stored.SecretAccessKey) +} + +func TestBackupService_S3ConfigFailsFastWhenVersionedSecretCannotBeDecrypted(t *testing.T) { + repo := newMockSettingRepo() + require.NoError(t, repo.Set(context.Background(), settingKeyBackupS3Config, + `{"bucket":"my-bucket","access_key_id":"AKID","secret_access_key":"enc:v1:broken"}`)) + svc := newTestBackupService(repo, &mockDumper{}, newMockObjectStore()) + + _, err := svc.loadS3Config(context.Background()) + + require.ErrorIs(t, err, ErrBackupS3ConfigCorrupt) +} + +func TestBackupService_UpdateS3ConfigRejectsStorageLocationChangesWithRecords(t *testing.T) { + base := BackupS3Config{ + Endpoint: "https://objects.example.com", + Region: "auto", + Bucket: "archive-bucket", + AccessKeyID: "AKID", + SecretAccessKey: "secret", + Prefix: "backups", + ForcePathStyle: false, + } + tests := []struct { + name string + mutate func(*BackupS3Config) + }{ + {name: "endpoint", mutate: func(cfg *BackupS3Config) { cfg.Endpoint = "https://other.example.com" }}, + {name: "region", mutate: func(cfg *BackupS3Config) { cfg.Region = "us-east-1" }}, + {name: "bucket", mutate: func(cfg *BackupS3Config) { cfg.Bucket = "other-bucket" }}, + {name: "prefix", mutate: func(cfg *BackupS3Config) { cfg.Prefix = "other-prefix" }}, + {name: "path_style", mutate: func(cfg *BackupS3Config) { cfg.ForcePathStyle = true }}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + repo := newMockSettingRepo() + store := newMockObjectStore() + svc := newTestBackupService(repo, &mockDumper{}, store) + _, err := svc.UpdateS3Config(context.Background(), base) + require.NoError(t, err) + seedBackupRecord(t, svc, store, BackupRecord{ + ID: "existing", + Status: "completed", + BackupType: "usage_logs_archive", + S3Key: "backups/2026/07/18/existing.gz", + StartedAt: time.Now().Format(time.RFC3339), + }) + + updated := base + updated.SecretAccessKey = "" + tt.mutate(&updated) + _, err = svc.UpdateS3Config(context.Background(), updated) + require.Error(t, err) + require.Equal(t, "BACKUP_STORAGE_LOCATION_IN_USE", infraerrors.Reason(err)) + }) + } +} + +func TestBackupService_UpdateS3ConfigAllowsCredentialRotationWithRecords(t *testing.T) { + repo := newMockSettingRepo() + store := newMockObjectStore() + svc := newTestBackupService(repo, &mockDumper{}, store) + base := BackupS3Config{ + Endpoint: "https://objects.example.com/", + Region: "auto", + Bucket: "archive-bucket", + AccessKeyID: "OLD-AKID", + SecretAccessKey: "old-secret", + Prefix: "backups", + } + _, err := svc.UpdateS3Config(context.Background(), base) + require.NoError(t, err) + seedBackupRecord(t, svc, store, BackupRecord{ + ID: "existing", + Status: "completed", + BackupType: "ops_error_logs_archive", + S3Key: "backups/2026/07/18/existing.gz", + StartedAt: time.Now().Format(time.RFC3339), + }) + + updated := base + updated.Endpoint = "https://objects.example.com" + updated.Prefix = "backups/" + updated.AccessKeyID = "NEW-AKID" + updated.SecretAccessKey = "new-secret" + _, err = svc.UpdateS3Config(context.Background(), updated) + require.NoError(t, err) + stored, err := svc.loadS3Config(context.Background()) + require.NoError(t, err) + require.Equal(t, "NEW-AKID", stored.AccessKeyID) + require.Equal(t, "new-secret", stored.SecretAccessKey) +} + +func TestBackupService_UpdateS3ConfigRejectsNewLocationWhenConfigWasRemoved(t *testing.T) { + repo := newMockSettingRepo() + store := newMockObjectStore() + svc := newTestBackupService(repo, &mockDumper{}, store) + seedBackupRecord(t, svc, store, BackupRecord{ + ID: "orphaned-location", + Status: "completed", + BackupType: "postgres", + S3Key: "backups/orphaned", + StartedAt: time.Now().Format(time.RFC3339), + }) + + _, err := svc.UpdateS3Config(context.Background(), BackupS3Config{ + Endpoint: "https://objects.example.com", + Region: "auto", + Bucket: "archive-bucket", + AccessKeyID: "AKID", + SecretAccessKey: "secret", + Prefix: "backups", + }) + require.Error(t, err) + require.Equal(t, "BACKUP_STORAGE_LOCATION_IN_USE", infraerrors.Reason(err)) } func TestBackupService_SaveRecordConcurrency(t *testing.T) { @@ -315,6 +526,27 @@ func TestBackupService_SaveRecordConcurrency(t *testing.T) { require.Len(t, records, n) } +func TestBackupService_SaveRecordRetainsObjectReferences(t *testing.T) { + repo := newMockSettingRepo() + svc := newTestBackupService(repo, &mockDumper{}, newMockObjectStore()) + + const recordCount = 101 + for i := 0; i < recordCount; i++ { + require.NoError(t, svc.saveRecord(context.Background(), &BackupRecord{ + ID: fmt.Sprintf("retained-%d", i), + Status: "completed", + S3Key: fmt.Sprintf("backups/retained-%d", i), + StartedAt: time.Now().Format(time.RFC3339), + })) + } + + records, err := svc.loadRecords(context.Background()) + require.NoError(t, err) + require.Len(t, records, recordCount) + require.Equal(t, "retained-0", records[0].ID) + require.Equal(t, "retained-100", records[len(records)-1].ID) +} + func TestBackupService_LoadRecords_Empty(t *testing.T) { repo := newMockSettingRepo() svc := newTestBackupService(repo, &mockDumper{}, newMockObjectStore()) @@ -334,6 +566,44 @@ func TestBackupService_LoadRecords_Corrupted(t *testing.T) { require.Nil(t, records) } +func TestBackupService_LoadRecords_EmptyValueIsCorrupt(t *testing.T) { + repo := newMockSettingRepo() + require.NoError(t, repo.Set(context.Background(), settingKeyBackupRecords, "")) + svc := newTestBackupService(repo, &mockDumper{}, newMockObjectStore()) + + records, err := svc.loadRecords(context.Background()) + require.ErrorIs(t, err, ErrBackupRecordsCorrupt) + require.Nil(t, records) +} + +func TestBackupService_LoadRecords_RepositoryErrorPropagates(t *testing.T) { + repo := newMockSettingRepo() + repoErr := fmt.Errorf("temporary setting repository failure") + repo.getValueErrors[settingKeyBackupRecords] = repoErr + svc := newTestBackupService(repo, &mockDumper{}, newMockObjectStore()) + + records, err := svc.loadRecords(context.Background()) + require.ErrorIs(t, err, repoErr) + require.Nil(t, records) +} + +func TestBackupService_SaveRecordLoadFailureDoesNotOverwriteRecords(t *testing.T) { + repo := newMockSettingRepo() + original := []BackupRecord{{ID: "original", Status: "completed", S3Key: "backups/original"}} + raw, err := json.Marshal(original) + require.NoError(t, err) + require.NoError(t, repo.Set(context.Background(), settingKeyBackupRecords, string(raw))) + repoErr := fmt.Errorf("temporary setting repository failure") + repo.getValueErrors[settingKeyBackupRecords] = repoErr + svc := newTestBackupService(repo, &mockDumper{}, newMockObjectStore()) + + err = svc.saveRecord(context.Background(), &BackupRecord{ID: "new", Status: "completed"}) + require.ErrorIs(t, err, repoErr) + repo.mu.Lock() + require.JSONEq(t, string(raw), repo.data[settingKeyBackupRecords]) + repo.mu.Unlock() +} + func TestBackupService_CreateBackup_Streaming(t *testing.T) { repo := newMockSettingRepo() seedS3Config(t, repo) @@ -388,6 +658,86 @@ func TestBackupServiceCreateUsageLogsArchive(t *testing.T) { require.Equal(t, "{\"id\":1}\n", string(plain)) } +func TestBackupService_CreateDataArchiveSaveFailureDeletesUploadedObject(t *testing.T) { + repo := newMockSettingRepo() + seedS3Config(t, repo) + persistErr := fmt.Errorf("backup record storage unavailable") + repo.setErrors[settingKeyBackupRecords] = persistErr + store := newMockObjectStore() + svc := newTestBackupService(repo, &mockDumper{}, store) + + record, err := svc.CreateDataArchive(context.Background(), DataArchiveInput{ + Stream: io.NopCloser(strings.NewReader("{\"id\":1}\n")), + FileName: "usage_logs.ndjson.gz", + BackupType: "usage_logs_archive", + TriggeredBy: "usage_cleanup_auto", + ExpireDays: 180, + }) + require.ErrorIs(t, err, persistErr) + require.NotNil(t, record) + require.Equal(t, "failed", record.Status) + store.mu.Lock() + require.Empty(t, store.objects) + require.Equal(t, []string{record.S3Key}, store.deleteCalls) + store.mu.Unlock() +} + +func TestBackupService_CreateDataArchiveCompensationFailureJoinsErrors(t *testing.T) { + repo := newMockSettingRepo() + seedS3Config(t, repo) + persistErr := fmt.Errorf("backup record storage unavailable") + repo.setErrors[settingKeyBackupRecords] = persistErr + compensationErr := fmt.Errorf("object storage delete unavailable") + store := newMockObjectStore() + store.deleteErr = compensationErr + svc := newTestBackupService(repo, &mockDumper{}, store) + + record, err := svc.CreateDataArchive(context.Background(), DataArchiveInput{ + Stream: io.NopCloser(strings.NewReader("{\"id\":1}\n")), + FileName: "usage_logs.ndjson.gz", + BackupType: "usage_logs_archive", + TriggeredBy: "usage_cleanup_auto", + ExpireDays: 180, + }) + require.ErrorIs(t, err, persistErr) + require.ErrorIs(t, err, compensationErr) + require.NotNil(t, record) + require.Equal(t, "failed", record.Status) + store.mu.Lock() + require.Contains(t, store.objects, record.S3Key) + require.Equal(t, []string{record.S3Key}, store.deleteCalls) + store.mu.Unlock() +} + +func TestBackupService_CreateDataArchiveDoesNotInlineCleanupExpiredObjects(t *testing.T) { + repo := newMockSettingRepo() + seedS3Config(t, repo) + store := newMockObjectStore() + svc := newTestBackupService(repo, &mockDumper{}, store) + now := time.Now().UTC().Truncate(time.Second) + seedBackupRecord(t, svc, store, BackupRecord{ + ID: "expired-existing", + Status: "completed", + BackupType: "usage_logs_archive", + S3Key: "expired/existing", + StartedAt: now.Add(-48 * time.Hour).Format(time.RFC3339), + ExpiresAt: now.Add(-time.Hour).Format(time.RFC3339), + }) + + _, err := svc.CreateDataArchive(context.Background(), DataArchiveInput{ + Stream: io.NopCloser(strings.NewReader("{\"id\":2}\n")), + FileName: "new.ndjson.gz", + BackupType: "usage_logs_archive", + TriggeredBy: "usage_cleanup_auto", + ExpireDays: 180, + }) + require.NoError(t, err) + store.mu.Lock() + require.Empty(t, store.deleteCalls) + require.Contains(t, store.objects, "expired/existing") + store.mu.Unlock() +} + func TestBackupServiceRestoreRejectsUsageLogsArchive(t *testing.T) { repo := newMockSettingRepo() seedS3Config(t, repo) @@ -445,6 +795,40 @@ func TestBackupService_CreateBackup_ConcurrentBlocked(t *testing.T) { require.ErrorIs(t, err, ErrBackupInProgress) } +func TestBackupService_BackupAndArchiveRejectedDuringRestore(t *testing.T) { + repo := newMockSettingRepo() + svc := newTestBackupService(repo, &mockDumper{}, newMockObjectStore()) + svc.opMu.Lock() + svc.restoring = true + svc.opMu.Unlock() + + _, err := svc.CreateBackup(context.Background(), "manual", 14) + require.ErrorIs(t, err, ErrRestoreInProgress) + _, err = svc.StartBackup(context.Background(), "manual", 14) + require.ErrorIs(t, err, ErrRestoreInProgress) + _, err = svc.CreateDataArchive(context.Background(), DataArchiveInput{ + Stream: io.NopCloser(strings.NewReader("{}\n")), + FileName: "archive.ndjson.gz", + BackupType: "usage_logs_archive", + }) + require.ErrorIs(t, err, ErrRestoreInProgress) + svc.wg.Wait() +} + +func TestBackupService_RestoreRejectedDuringBackup(t *testing.T) { + repo := newMockSettingRepo() + svc := newTestBackupService(repo, &mockDumper{}, newMockObjectStore()) + svc.opMu.Lock() + svc.backingUp = true + svc.opMu.Unlock() + + err := svc.RestoreBackup(context.Background(), "backup-id") + require.ErrorIs(t, err, ErrBackupInProgress) + _, err = svc.StartRestore(context.Background(), "backup-id") + require.ErrorIs(t, err, ErrBackupInProgress) + svc.wg.Wait() +} + func TestBackupService_RestoreBackup_Streaming(t *testing.T) { repo := newMockSettingRepo() seedS3Config(t, repo) @@ -593,6 +977,139 @@ func TestBackupService_Schedule_CronValidation(t *testing.T) { require.Error(t, err) } +func TestBackupService_StartRegistersExpirationCleanupWhenFullBackupDisabled(t *testing.T) { + repo := newMockSettingRepo() + seedS3Config(t, repo) + scheduleData, err := json.Marshal(BackupScheduleConfig{Enabled: false}) + require.NoError(t, err) + require.NoError(t, repo.Set(context.Background(), settingKeyBackupSchedule, string(scheduleData))) + store := newMockObjectStore() + svc := newTestBackupService(repo, &mockDumper{}, store) + now := time.Now().UTC().Truncate(time.Second) + seedBackupRecord(t, svc, store, BackupRecord{ + ID: "expired-before-start", + Status: "completed", + BackupType: "usage_logs_archive", + S3Key: "expired/before-start", + StartedAt: now.Add(-2 * time.Hour).Format(time.RFC3339), + ExpiresAt: now.Add(-time.Hour).Format(time.RFC3339), + }) + + svc.Start() + defer svc.Stop() + + svc.cronMu.Lock() + cleanupEntryID := svc.cronCleanupEntryID + backupEntryID := svc.cronEntryID + entries := svc.cronSched.Entries() + svc.cronMu.Unlock() + require.NotZero(t, cleanupEntryID) + require.Zero(t, backupEntryID) + require.Len(t, entries, 1) + require.Equal(t, cleanupEntryID, entries[0].ID) + require.Equal(t, 3, entries[0].Next.Hour()) + require.Equal(t, 30, entries[0].Next.Minute()) + + // Start 只注册固定低峰任务，不立即执行删除。 + _, err = svc.GetBackupRecord(context.Background(), "expired-before-start") + require.NoError(t, err) + store.mu.Lock() + require.Empty(t, store.deleteCalls) + store.mu.Unlock() +} + +func TestBackupService_StopDisablesExpirationCleanupLifecycle(t *testing.T) { + repo := newMockSettingRepo() + seedS3Config(t, repo) + store := newMockObjectStore() + svc := newTestBackupService(repo, &mockDumper{}, store) + now := time.Now().UTC().Truncate(time.Second) + seedBackupRecord(t, svc, store, BackupRecord{ + ID: "expired-after-stop", + Status: "completed", + BackupType: "ops_error_logs_archive", + S3Key: "expired/after-stop", + StartedAt: now.Add(-2 * time.Hour).Format(time.RFC3339), + ExpiresAt: now.Add(-time.Hour).Format(time.RFC3339), + }) + + svc.Start() + svc.Stop() + require.True(t, svc.shuttingDown.Load()) + svc.cronMu.Lock() + require.Zero(t, svc.cronCleanupEntryID) + require.Zero(t, svc.cronEntryID) + svc.cronMu.Unlock() + + // 即使停止后出现迟到的定时回调，也不得进入清理或增加 WaitGroup。 + svc.runScheduledExpirationCleanup() + svc.wg.Wait() + _, err := svc.GetBackupRecord(context.Background(), "expired-after-stop") + require.NoError(t, err) + store.mu.Lock() + require.Empty(t, store.deleteCalls) + store.mu.Unlock() +} + +func TestScheduledExpirationCleanup_BusyRetriesAreBounded(t *testing.T) { + repo := newMockSettingRepo() + seedS3Config(t, repo) + store := newMockObjectStore() + svc := newTestBackupService(repo, &mockDumper{}, store) + now := time.Now().UTC().Truncate(time.Second) + seedBackupRecord(t, svc, store, BackupRecord{ + ID: "expired-during-backup", + Status: "completed", + BackupType: "ops_system_logs_archive", + S3Key: "expired/during-backup", + StartedAt: now.Add(-2 * time.Hour).Format(time.RFC3339), + ExpiresAt: now.Add(-time.Hour).Format(time.RFC3339), + }) + svc.opMu.Lock() + svc.backingUp = true + svc.opMu.Unlock() + + err := svc.cleanupExpiredBackupsWithRetry(context.Background(), 3, time.Millisecond) + require.ErrorIs(t, err, ErrBackupInProgress) + _, err = svc.GetBackupRecord(context.Background(), "expired-during-backup") + require.NoError(t, err) + store.mu.Lock() + require.Empty(t, store.deleteCalls) + store.mu.Unlock() +} + +func TestScheduledExpirationCleanup_RetriesAfterBackupFinishes(t *testing.T) { + repo := newMockSettingRepo() + seedS3Config(t, repo) + store := newMockObjectStore() + svc := newTestBackupService(repo, &mockDumper{}, store) + now := time.Now().UTC().Truncate(time.Second) + seedBackupRecord(t, svc, store, BackupRecord{ + ID: "expired-after-busy", + Status: "completed", + BackupType: "ops_system_logs_archive", + S3Key: "expired/after-busy", + StartedAt: now.Add(-2 * time.Hour).Format(time.RFC3339), + ExpiresAt: now.Add(-time.Hour).Format(time.RFC3339), + }) + svc.opMu.Lock() + svc.backingUp = true + svc.opMu.Unlock() + time.AfterFunc(5*time.Millisecond, func() { + svc.opMu.Lock() + svc.backingUp = false + svc.opMu.Unlock() + }) + + err := svc.cleanupExpiredBackupsWithRetry(context.Background(), 2, 20*time.Millisecond) + require.NoError(t, err) + _, err = svc.GetBackupRecord(context.Background(), "expired-after-busy") + require.ErrorIs(t, err, ErrBackupNotFound) + store.mu.Lock() + require.Equal(t, []string{"expired/after-busy"}, store.deleteCalls) + store.mu.Unlock() +} + func TestBackupService_LoadS3Config_Corrupted(t *testing.T) { repo := newMockSettingRepo() _ = repo.Set(context.Background(), settingKeyBackupS3Config, "not json!!!!") @@ -603,6 +1120,19 @@ func TestBackupService_LoadS3Config_Corrupted(t *testing.T) { require.Nil(t, cfg) } +func TestBackupService_LoadS3Config_RepositoryErrorPropagates(t *testing.T) { + repo := newMockSettingRepo() + repoErr := fmt.Errorf("temporary setting repository failure") + repo.getValueErrors[settingKeyBackupS3Config] = repoErr + svc := newTestBackupService(repo, &mockDumper{}, newMockObjectStore()) + + cfg, err := svc.loadS3Config(context.Background()) + require.ErrorIs(t, err, repoErr) + require.Nil(t, cfg) + err = svc.TestS3Connection(context.Background(), BackupS3Config{Bucket: "bucket", AccessKeyID: "AKID"}) + require.ErrorIs(t, err, repoErr) +} + // ─── Async Backup Tests ─── func TestStartBackup_ReturnsImmediately(t *testing.T) { @@ -698,33 +1228,180 @@ func TestGracefulShutdown(t *testing.T) { store := newMockObjectStore() svc := newTestBackupService(repo, dumper, store) - _, err := svc.StartBackup(context.Background(), "manual", 14) + record, err := svc.StartBackup(context.Background(), "manual", 14) require.NoError(t, err) - // Stop 应该等待备份完成 + // Stop 先取消后台 context，再有界等待 goroutine 自行收口；无需手工释放 dumper。 done := make(chan struct{}) go func() { svc.Stop() close(done) }() - // 短暂等待确认 Stop 还在等待 select { case <-done: - t.Fatal("Stop returned before backup finished") - case <-time.After(100 * time.Millisecond): - // 预期：Stop 还在等待 + case <-time.After(2 * time.Second): + close(dumper.blockCh) + t.Fatal("Stop did not cancel and finish the active backup within the bounded wait") } - - // 释放备份 close(dumper.blockCh) + final, err := svc.GetBackupRecord(context.Background(), record.ID) + require.NoError(t, err) + require.Equal(t, "failed", final.Status) + require.Contains(t, final.ErrorMsg, "context canceled") +} - // 现在 Stop 应该完成 +func TestBackupService_StopRejectsNewOperations(t *testing.T) { + repo := newMockSettingRepo() + seedS3Config(t, repo) + store := newMockObjectStore() + svc := newTestBackupService(repo, &mockDumper{dumpData: []byte("data")}, store) + seedBackupRecord(t, svc, store, BackupRecord{ + ID: "existing", + Status: "completed", + BackupType: "postgres", + S3Key: "backups/existing", + StartedAt: time.Now().Format(time.RFC3339), + }) + svc.Stop() + + assertShuttingDown := func(err error) { + t.Helper() + require.Error(t, err) + require.Equal(t, "SERVER_SHUTTING_DOWN", infraerrors.Reason(err)) + } + _, err := svc.CreateBackup(context.Background(), "manual", 14) + assertShuttingDown(err) + _, err = svc.StartBackup(context.Background(), "manual", 14) + assertShuttingDown(err) + _, err = svc.CreateDataArchive(context.Background(), DataArchiveInput{ + Stream: io.NopCloser(strings.NewReader("{}\n")), + FileName: "archive.ndjson.gz", + BackupType: "usage_logs_archive", + }) + assertShuttingDown(err) + err = svc.RestoreBackup(context.Background(), "existing") + assertShuttingDown(err) + _, err = svc.StartRestore(context.Background(), "existing") + assertShuttingDown(err) + err = svc.DeleteBackup(context.Background(), "existing") + assertShuttingDown(err) + _, err = svc.UpdateS3Config(context.Background(), BackupS3Config{}) + assertShuttingDown(err) +} + +func TestBackupService_StopSerializesWaitGroupRegistration(t *testing.T) { + svc := newTestBackupService(newMockSettingRepo(), &mockDumper{}, newMockObjectStore()) + start := make(chan struct{}) + var callers sync.WaitGroup + for i := 0; i < 100; i++ { + callers.Add(1) + go func() { + defer callers.Done() + <-start + if svc.tryBeginRun() { + time.Sleep(time.Microsecond) + svc.endRun() + } + }() + } + stopDone := make(chan struct{}) + go func() { + <-start + svc.Stop() + close(stopDone) + }() + close(start) + callers.Wait() select { - case <-done: - // 预期 - case <-time.After(5 * time.Second): - t.Fatal("Stop did not return after backup finished") + case <-stopDone: + case <-time.After(2 * time.Second): + t.Fatal("Stop did not finish after registered operations ended") + } + require.False(t, svc.tryBeginRun()) +} + +func TestStartBackup_InitialRecordFailureBalancesLifecycle(t *testing.T) { + repo := newMockSettingRepo() + seedS3Config(t, repo) + persistErr := fmt.Errorf("backup record storage unavailable") + repo.setErrors[settingKeyBackupRecords] = persistErr + svc := newTestBackupService(repo, &mockDumper{dumpData: []byte("data")}, newMockObjectStore()) + + _, err := svc.StartBackup(context.Background(), "manual", 14) + require.ErrorIs(t, err, persistErr) + svc.opMu.Lock() + require.False(t, svc.backingUp) + svc.opMu.Unlock() + waitDone := make(chan struct{}) + go func() { + svc.wg.Wait() + close(waitDone) + }() + select { + case <-waitDone: + case <-time.After(time.Second): + t.Fatal("StartBackup leaked its lifecycle registration after initialization failure") + } +} + +func TestStartBackup_FinalRecordFailureCompensatesUploadedObject(t *testing.T) { + repo := newMockSettingRepo() + seedS3Config(t, repo) + persistErr := fmt.Errorf("backup record storage unavailable") + store := newMockObjectStore() + store.uploadHook = func() { + repo.mu.Lock() + repo.setErrors[settingKeyBackupRecords] = persistErr + repo.mu.Unlock() + } + svc := newTestBackupService(repo, &mockDumper{dumpData: []byte("data")}, store) + + record, err := svc.StartBackup(context.Background(), "manual", 14) + require.NoError(t, err) + svc.wg.Wait() + repo.mu.Lock() + delete(repo.setErrors, settingKeyBackupRecords) + repo.mu.Unlock() + final, err := svc.GetBackupRecord(context.Background(), record.ID) + require.NoError(t, err) + // 初始 running 记录存在；最终状态保存失败时不能伪装为 completed，且上传对象已补偿删除。 + require.Equal(t, "running", final.Status) + store.mu.Lock() + require.Empty(t, store.objects) + require.Equal(t, []string{record.S3Key}, store.deleteCalls) + store.mu.Unlock() +} + +func TestStartRestore_StatusSaveFailureBalancesLifecycle(t *testing.T) { + repo := newMockSettingRepo() + seedS3Config(t, repo) + store := newMockObjectStore() + svc := newTestBackupService(repo, &mockDumper{}, store) + seedBackupRecord(t, svc, store, BackupRecord{ + ID: "restore-source", + Status: "completed", + BackupType: "postgres", + S3Key: "backups/restore-source", + StartedAt: time.Now().Format(time.RFC3339), + }) + persistErr := fmt.Errorf("backup record storage unavailable") + repo.setErrors[settingKeyBackupRecords] = persistErr + + _, err := svc.StartRestore(context.Background(), "restore-source") + require.ErrorIs(t, err, persistErr) + svc.opMu.Lock() + require.False(t, svc.restoring) + svc.opMu.Unlock() + waitDone := make(chan struct{}) + go func() { + svc.wg.Wait() + close(waitDone) + }() + select { + case <-waitDone: + case <-time.After(time.Second): + t.Fatal("StartRestore leaked its lifecycle registration after initialization failure") } } @@ -753,3 +1430,198 @@ func TestStartRestore_Async(t *testing.T) { require.NoError(t, err) require.Equal(t, "completed", final.RestoreStatus) } + +func TestCleanupExpiredBackups_ArchiveTypesOnly(t *testing.T) { + repo := newMockSettingRepo() + seedS3Config(t, repo) + store := newMockObjectStore() + svc := newTestBackupService(repo, &mockDumper{}, store) + now := time.Now().UTC().Truncate(time.Second) + + for _, record := range []BackupRecord{ + {ID: "expired-usage", Status: "completed", BackupType: "usage_logs_archive", S3Key: "expired/usage", StartedAt: now.Add(-48 * time.Hour).Format(time.RFC3339), ExpiresAt: now.Add(-time.Hour).Format(time.RFC3339)}, + {ID: "expired-ops-system", Status: "completed", BackupType: "ops_system_logs_archive", S3Key: "expired/ops-system", StartedAt: now.Add(-48 * time.Hour).Format(time.RFC3339), ExpiresAt: now.Add(-time.Hour).Format(time.RFC3339)}, + {ID: "expired-postgres", Status: "completed", BackupType: "postgres", S3Key: "expired/postgres", StartedAt: now.Add(-48 * time.Hour).Format(time.RFC3339), ExpiresAt: now.Add(-time.Hour).Format(time.RFC3339)}, + {ID: "future-ops-error", Status: "completed", BackupType: "ops_error_logs_archive", S3Key: "future/ops-error", StartedAt: now.Format(time.RFC3339), ExpiresAt: now.Add(time.Hour).Format(time.RFC3339)}, + {ID: "never-expires", Status: "completed", BackupType: "usage_logs_archive", S3Key: "future/never", StartedAt: now.Format(time.RFC3339)}, + } { + seedBackupRecord(t, svc, store, record) + } + + require.NoError(t, svc.cleanupExpiredBackups(context.Background())) + for _, id := range []string{"expired-usage", "expired-ops-system"} { + _, err := svc.GetBackupRecord(context.Background(), id) + require.ErrorIs(t, err, ErrBackupNotFound) + } + for _, id := range []string{"expired-postgres", "future-ops-error", "never-expires"} { + _, err := svc.GetBackupRecord(context.Background(), id) + require.NoError(t, err) + } + store.mu.Lock() + require.ElementsMatch(t, []string{"expired/usage", "expired/ops-system"}, store.deleteCalls) + require.Len(t, store.objects, 3) + store.mu.Unlock() +} + +func TestCleanupExpiredBackups_FailuresRemainAndAreReported(t *testing.T) { + repo := newMockSettingRepo() + seedS3Config(t, repo) + store := newMockObjectStore() + svc := newTestBackupService(repo, &mockDumper{}, store) + now := time.Now().UTC().Truncate(time.Second) + expired := now.Add(-time.Hour).Format(time.RFC3339) + + for _, record := range []BackupRecord{ + {ID: "deleted", Status: "completed", BackupType: "usage_logs_archive", S3Key: "expired/deleted", StartedAt: now.Add(-time.Hour).Format(time.RFC3339), ExpiresAt: expired}, + {ID: "delete-failed", Status: "completed", BackupType: "ops_error_logs_archive", S3Key: "expired/failure", StartedAt: now.Add(-time.Hour).Format(time.RFC3339), ExpiresAt: expired}, + {ID: "invalid-expiry", Status: "completed", BackupType: "usage_logs_archive", S3Key: "expired/invalid", StartedAt: now.Add(-time.Hour).Format(time.RFC3339), ExpiresAt: "invalid"}, + {ID: "empty-key", Status: "completed", BackupType: "ops_system_logs_archive", StartedAt: now.Add(-time.Hour).Format(time.RFC3339), ExpiresAt: expired}, + } { + seedBackupRecord(t, svc, store, record) + } + store.mu.Lock() + store.deleteErrors["expired/failure"] = fmt.Errorf("S3 unavailable") + store.mu.Unlock() + + err := svc.cleanupExpiredBackups(context.Background()) + require.Error(t, err) + require.Contains(t, err.Error(), "delete-failed") + require.Contains(t, err.Error(), "invalid-expiry") + require.Contains(t, err.Error(), "empty-key") + _, err = svc.GetBackupRecord(context.Background(), "deleted") + require.ErrorIs(t, err, ErrBackupNotFound) + for _, id := range []string{"delete-failed", "invalid-expiry", "empty-key"} { + _, err = svc.GetBackupRecord(context.Background(), id) + require.NoError(t, err) + } +} + +func TestCleanupExpiredBackups_DeletesAtMostOneBatch(t *testing.T) { + repo := newMockSettingRepo() + seedS3Config(t, repo) + store := newMockObjectStore() + svc := newTestBackupService(repo, &mockDumper{}, store) + now := time.Now().UTC().Truncate(time.Second) + for i := 0; i < backupExpirationCleanupBatchSize+1; i++ { + seedBackupRecord(t, svc, store, BackupRecord{ + ID: fmt.Sprintf("expired-%d", i), + Status: "completed", + BackupType: "usage_logs_archive", + S3Key: fmt.Sprintf("expired/%d", i), + StartedAt: now.Add(-48 * time.Hour).Format(time.RFC3339), + ExpiresAt: now.Add(-time.Hour).Format(time.RFC3339), + }) + } + + require.NoError(t, svc.cleanupExpiredBackups(context.Background())) + store.mu.Lock() + require.Len(t, store.deleteCalls, backupExpirationCleanupBatchSize) + store.mu.Unlock() + records, err := svc.loadRecords(context.Background()) + require.NoError(t, err) + require.Len(t, records, 1) +} + +func TestCleanupOldBackups_PostgresOnly(t *testing.T) { + repo := newMockSettingRepo() + seedS3Config(t, repo) + store := newMockObjectStore() + svc := newTestBackupService(repo, &mockDumper{}, store) + now := time.Now().UTC().Truncate(time.Second) + + for _, record := range []BackupRecord{ + {ID: "pg-new", Status: "completed", BackupType: "postgres", S3Key: "pg/new", StartedAt: now.Format(time.RFC3339)}, + {ID: "pg-old", Status: "completed", BackupType: "postgres", S3Key: "pg/old", StartedAt: now.AddDate(0, 0, -60).Format(time.RFC3339)}, + {ID: "legacy-pg-old", Status: "completed", S3Key: "pg/legacy", StartedAt: now.AddDate(0, 0, -90).Format(time.RFC3339)}, + {ID: "usage-old", Status: "completed", BackupType: "usage_logs_archive", S3Key: "archive/usage", StartedAt: now.AddDate(0, 0, -90).Format(time.RFC3339)}, + {ID: "ops-old", Status: "completed", BackupType: "ops_error_logs_archive", S3Key: "archive/ops", StartedAt: now.AddDate(0, 0, -90).Format(time.RFC3339)}, + } { + seedBackupRecord(t, svc, store, record) + } + + require.NoError(t, svc.cleanupOldBackups(context.Background(), &BackupScheduleConfig{RetainDays: 30, RetainCount: 1})) + for _, id := range []string{"pg-old", "legacy-pg-old"} { + _, err := svc.GetBackupRecord(context.Background(), id) + require.ErrorIs(t, err, ErrBackupNotFound) + } + for _, id := range []string{"pg-new", "usage-old", "ops-old"} { + _, err := svc.GetBackupRecord(context.Background(), id) + require.NoError(t, err) + } + store.mu.Lock() + require.ElementsMatch(t, []string{"pg/old", "pg/legacy"}, store.deleteCalls) + store.mu.Unlock() +} + +func TestCleanupOldBackups_AllPostgresBackupsExpiredKeepsLatest(t *testing.T) { + repo := newMockSettingRepo() + seedS3Config(t, repo) + store := newMockObjectStore() + svc := newTestBackupService(repo, &mockDumper{}, store) + now := time.Now().UTC().Truncate(time.Second) + for i := 0; i < 3; i++ { + seedBackupRecord(t, svc, store, BackupRecord{ + ID: fmt.Sprintf("pg-%d", i), + Status: "completed", + BackupType: "postgres", + S3Key: fmt.Sprintf("pg/%d", i), + StartedAt: now.AddDate(0, 0, -(90 - i)).Format(time.RFC3339), + }) + } + + require.NoError(t, svc.cleanupOldBackups(context.Background(), &BackupScheduleConfig{RetainDays: 30, RetainCount: 0})) + _, err := svc.GetBackupRecord(context.Background(), "pg-2") + require.NoError(t, err) + for _, id := range []string{"pg-0", "pg-1"} { + _, err = svc.GetBackupRecord(context.Background(), id) + require.ErrorIs(t, err, ErrBackupNotFound) + } + store.mu.Lock() + require.ElementsMatch(t, []string{"pg/0", "pg/1"}, store.deleteCalls) + store.mu.Unlock() +} + +func TestCleanupExpiredBackups_BlocksConcurrentBackup(t *testing.T) { + repo := newMockSettingRepo() + seedS3Config(t, repo) + store := newMockObjectStore() + deleteBlock := make(chan struct{}) + store.deleteBlock = deleteBlock + store.deleteStarted = make(chan struct{}) + svc := newTestBackupService(repo, &mockDumper{dumpData: []byte("data")}, store) + now := time.Now().UTC().Truncate(time.Second) + seedBackupRecord(t, svc, store, BackupRecord{ + ID: "expired", Status: "completed", BackupType: "usage_logs_archive", S3Key: "expired/object", + StartedAt: now.Add(-time.Hour).Format(time.RFC3339), ExpiresAt: now.Add(-time.Minute).Format(time.RFC3339), + }) + + cleanupDone := make(chan error, 1) + go func() { cleanupDone <- svc.cleanupExpiredBackups(context.Background()) }() + select { + case <-store.deleteStarted: + case <-time.After(2 * time.Second): + t.Fatal("expiration cleanup did not reach object deletion") + } + _, err := svc.CreateBackup(context.Background(), "manual", 14) + require.ErrorIs(t, err, ErrBackupInProgress) + close(deleteBlock) + require.NoError(t, <-cleanupDone) +} + +func TestDeleteBackup_DeleteFailureKeepsRecord(t *testing.T) { + repo := newMockSettingRepo() + seedS3Config(t, repo) + store := newMockObjectStore() + svc := newTestBackupService(repo, &mockDumper{}, store) + seedBackupRecord(t, svc, store, BackupRecord{ + ID: "keep-on-error", Status: "completed", BackupType: "postgres", S3Key: "delete/failure", StartedAt: time.Now().Format(time.RFC3339), + }) + store.mu.Lock() + store.deleteErrors["delete/failure"] = fmt.Errorf("S3 unavailable") + store.mu.Unlock() + + err := svc.DeleteBackup(context.Background(), "keep-on-error") + require.Error(t, err) + _, getErr := svc.GetBackupRecord(context.Background(), "keep-on-error") + require.NoError(t, getErr) +} diff --git a/backend/internal/service/billing_service.go b/backend/internal/service/billing_service.go index 6ce357c04..f0c677d3b 100644 --- a/backend/internal/service/billing_service.go +++ b/backend/internal/service/billing_service.go @@ -118,7 +118,8 @@ type UsageTokens struct { // CostBreakdown 费用明细 type CostBreakdown struct { - InputCost float64 + InputCost float64 // 文本输入费用，不含图片输入 + ImageInputCost float64 // 图片输入 token 费用 OutputCost float64 ImageOutputCost float64 CacheCreationCost float64 @@ -667,8 +668,8 @@ func (s *BillingService) computeTokenBreakdown( bd := &CostBreakdown{} textInputTokens := tokens.InputTokens + tokens.TextInputTokens - bd.InputCost = float64(textInputTokens)*inputPrice*inputTierMultiplier + - float64(tokens.ImageInputTokens)*imageInputPrice*imageInputTierMultiplier + bd.InputCost = float64(textInputTokens) * inputPrice * inputTierMultiplier + bd.ImageInputCost = float64(tokens.ImageInputTokens) * imageInputPrice * imageInputTierMultiplier // 分离图片输出 token 与文本输出 token textOutputTokens := tokens.OutputTokens - tokens.ImageOutputTokens @@ -695,7 +696,7 @@ func (s *BillingService) computeTokenBreakdown( bd.CacheReadCost = float64(textCacheReadTokens)*cacheReadPrice*cacheReadTierMultiplier + float64(imageCacheReadTokens)*imageCacheReadPrice*imageCacheReadTierMultiplier - bd.TotalCost = bd.InputCost + bd.OutputCost + bd.ImageOutputCost + + bd.TotalCost = bd.InputCost + bd.ImageInputCost + bd.OutputCost + bd.ImageOutputCost + bd.CacheCreationCost + bd.CacheReadCost bd.ActualCost = bd.TotalCost * rateMultiplier @@ -928,6 +929,7 @@ func (s *BillingService) CalculateCostWithLongContext(model string, tokens Usage // 合并成本 return &CostBreakdown{ InputCost: inRangeCost.InputCost + outRangeCost.InputCost, + ImageInputCost: inRangeCost.ImageInputCost + outRangeCost.ImageInputCost, OutputCost: inRangeCost.OutputCost, ImageOutputCost: inRangeCost.ImageOutputCost, CacheCreationCost: inRangeCost.CacheCreationCost, diff --git a/backend/internal/service/billing_service_unified_test.go b/backend/internal/service/billing_service_unified_test.go index 8c0f1df59..732d4d74e 100644 --- a/backend/internal/service/billing_service_unified_test.go +++ b/backend/internal/service/billing_service_unified_test.go @@ -220,11 +220,13 @@ func TestCalculateCostUnified_ImageTokenPrices(t *testing.T) { }) require.NoError(t, err) - expectedInput := 22*5e-6 + 10*8e-6 + expectedTextInput := 22 * 5e-6 + expectedImageInput := 10 * 8e-6 expectedCacheRead := 1*1.25e-6 + 3*2e-6 expectedImageOutput := 196 * 30e-6 - expectedTotal := expectedInput + expectedCacheRead + expectedImageOutput - require.InDelta(t, expectedInput, cost.InputCost, 1e-12) + expectedTotal := expectedTextInput + expectedImageInput + expectedCacheRead + expectedImageOutput + require.InDelta(t, expectedTextInput, cost.InputCost, 1e-12) + require.InDelta(t, expectedImageInput, cost.ImageInputCost, 1e-12) require.InDelta(t, 0.0, cost.OutputCost, 1e-12) require.InDelta(t, expectedImageOutput, cost.ImageOutputCost, 1e-12) require.InDelta(t, expectedCacheRead, cost.CacheReadCost, 1e-12) diff --git a/backend/internal/service/channel_monitor_checker.go b/backend/internal/service/channel_monitor_checker.go index d9b036cf4..405c1b5f3 100644 --- a/backend/internal/service/channel_monitor_checker.go +++ b/backend/internal/service/channel_monitor_checker.go @@ -151,7 +151,7 @@ func pingEndpointOrigin(ctx context.Context, endpoint string) *int { // - 拼出请求路径（含 model 占位） // - 序列化请求体 // - 构造鉴权头 -// - 从响应 JSON 中按 path 提取文本（gjson path） +// - 从响应 JSON 中提取文本（默认按 gjson path，需要时可自定义） // // 加新 provider 只需要在 providerAdapters 里增加一个条目，无需触碰 callProvider / validateProvider。 type providerAdapter struct { @@ -159,6 +159,7 @@ type providerAdapter struct { buildBody func(model, prompt string) ([]byte, error) buildHeaders func(apiKey string) map[string]string textPath string // gjson 提取响应文本的 path + extractText func([]byte) string } // providerAdapters 全部已支持的 provider。键值即 MonitorProvider* 字符串。 @@ -210,7 +211,7 @@ var providerAdapters = map[string]providerAdapter{ "anthropic-version": monitorAnthropicAPIVersion, } }, - textPath: "content.0.text", + extractText: extractAnthropicMonitorText, }, MonitorProviderGemini: { // Gemini 把 model 名写在 URL path 上：/v1beta/models/{model}:generateContent @@ -261,7 +262,34 @@ func callProvider(ctx context.Context, provider, endpoint, apiKey, model, prompt if err != nil { return "", "", status, err } - return gjson.GetBytes(respBytes, adapter.textPath).String(), string(respBytes), status, nil + return extractMonitorResponseText(adapter, respBytes), string(respBytes), status, nil +} + +func extractMonitorResponseText(adapter providerAdapter, respBytes []byte) string { + if adapter.extractText != nil { + return adapter.extractText(respBytes) + } + return gjson.GetBytes(respBytes, adapter.textPath).String() +} + +func extractAnthropicMonitorText(respBytes []byte) string { + content := gjson.GetBytes(respBytes, "content") + if !content.IsArray() { + return "" + } + + parts := make([]string, 0, 1) + content.ForEach(func(_, item gjson.Result) bool { + if item.Get("type").String() != "text" { + return true + } + text := strings.TrimSpace(item.Get("text").String()) + if text != "" { + parts = append(parts, text) + } + return true + }) + return strings.Join(parts, "\n") } // mergeHeaders 把用户自定义 headers 合并到 adapter 默认 headers 上。 diff --git a/backend/internal/service/channel_monitor_checker_body_test.go b/backend/internal/service/channel_monitor_checker_body_test.go index 323cf8b70..899669800 100644 --- a/backend/internal/service/channel_monitor_checker_body_test.go +++ b/backend/internal/service/channel_monitor_checker_body_test.go @@ -10,6 +10,8 @@ import ( "strings" "testing" "time" + + "github.com/stretchr/testify/require" ) // swapMonitorHTTPClient 临时替换 monitorHTTPClient 为不带 SSRF 校验的普通 client， @@ -171,3 +173,21 @@ func TestRunCheckForModel_ReplaceMode_EmptyResponseIsFailed(t *testing.T) { t.Errorf("failure message should hint replace-mode, got %q", res.Message) } } + +func TestExtractAnthropicMonitorText(t *testing.T) { + tests := []struct { + name string + body string + want string + }{ + {name: "text after thinking", body: `{"content":[{"type":"thinking","thinking":""},{"type":"text","text":"2"}]}`, want: "2"}, + {name: "multiple text blocks", body: `{"content":[{"type":"text","text":"answer"},{"type":"tool_use","name":"x"},{"type":"text","text":"2"}]}`, want: "answer\n2"}, + {name: "thinking only", body: `{"content":[{"type":"thinking","thinking":""}]}`, want: ""}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + require.Equal(t, tt.want, extractAnthropicMonitorText([]byte(tt.body))) + }) + } +} diff --git a/backend/internal/service/chatcompletions_anthropic_bridge.go b/backend/internal/service/chatcompletions_anthropic_bridge.go index 45c28ceb5..12b52cda9 100644 --- a/backend/internal/service/chatcompletions_anthropic_bridge.go +++ b/backend/internal/service/chatcompletions_anthropic_bridge.go @@ -1138,10 +1138,14 @@ func (s *OpenAIGatewayService) forwardAnthropicViaRawChatCompletions( upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(respBody))) if s.shouldFailoverOpenAIUpstreamResponse(resp.StatusCode, upstreamMsg, respBody) { appendOpsUpstreamError(c, OpsUpstreamErrorEvent{Platform: account.Platform, AccountID: account.ID, AccountName: account.Name, UpstreamStatusCode: resp.StatusCode, UpstreamRequestID: resp.Header.Get("x-request-id"), Kind: "failover", Message: upstreamMsg}) - if s.rateLimitService != nil { - s.rateLimitService.HandleUpstreamErrorForModel(ctx, account, upstreamModel, resp.StatusCode, resp.Header, respBody) - } - return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: respBody, RetryableOnSameAccount: shouldRetryOpenAIOnSamePoolAccount(account, resp.StatusCode, upstreamMsg, respBody)} + s.handleOpenAIAccountUpstreamErrorForModel(ctx, account, originalModel, resp.StatusCode, resp.Header, respBody) + return nil, newOpenAIUpstreamFailoverError( + resp.StatusCode, + resp.Header, + respBody, + upstreamMsg, + shouldRetryOpenAIOnSamePoolAccount(account, resp.StatusCode, upstreamMsg, respBody), + ) } return s.handleAnthropicErrorResponse(resp, c, account, billingModel) } diff --git a/backend/internal/service/domain_constants.go b/backend/internal/service/domain_constants.go index c92793919..6f611be6c 100644 --- a/backend/internal/service/domain_constants.go +++ b/backend/internal/service/domain_constants.go @@ -245,6 +245,7 @@ const ( SettingKeyWithdrawalManagementEnabled = "withdrawal_management_enabled" // 是否启用提现管理（默认开启） SettingKeyWithdrawalRateLimitWindowDays = "withdrawal_rate_limit_window_days" // 提现滚动频次限制窗口（天） SettingKeyWithdrawalRateLimitMax = "withdrawal_rate_limit_max" // 窗口内最多提现申请次数，0 表示不限制 + SettingKeyWithdrawalRateLimitExemptAmount = "withdrawal_rate_limit_exempt_amount" // 单笔提现严格超过该金额时免计频次，0 表示关闭豁免 SettingKeyContentModerationConfig = "content_moderation_config" // 内容审计配置（JSON） SettingKeyAccountShareCommentReviewEnabled = "account_share_comment_review_enabled" // 账号广场评论审核开关 SettingKeyAccountShareCommentReviewURL = "account_share_comment_review_url" // 账号广场评论审核模型 URL diff --git a/backend/internal/service/functional_module_switch_test.go b/backend/internal/service/functional_module_switch_test.go index 57464373e..80f438dda 100644 --- a/backend/internal/service/functional_module_switch_test.go +++ b/backend/internal/service/functional_module_switch_test.go @@ -24,7 +24,7 @@ func TestSettingServiceWithdrawalRateLimitDefaults(t *testing.T) { config, err := svc.GetWithdrawalRateLimitConfig(context.Background()) require.NoError(t, err) - require.Equal(t, WithdrawalRateLimitConfig{WindowDays: 1, MaxRequests: 0}, config) + require.Equal(t, WithdrawalRateLimitConfig{WindowDays: 1, MaxRequests: 0, ExemptAmount: 500}, config) } func TestSettingServiceWithdrawalRateLimitRejectsInvalidStoredValue(t *testing.T) { @@ -57,9 +57,10 @@ func TestWithdrawalServiceSubmitPassesConfiguredRateLimit(t *testing.T) { submitResult: &WithdrawalRequest{UserID: 1}, } settingSvc := newModuleSwitchSettingService(map[string]string{ - SettingKeyWithdrawalManagementEnabled: "true", - SettingKeyWithdrawalRateLimitWindowDays: "7", - SettingKeyWithdrawalRateLimitMax: "3", + SettingKeyWithdrawalManagementEnabled: "true", + SettingKeyWithdrawalRateLimitWindowDays: "7", + SettingKeyWithdrawalRateLimitMax: "3", + SettingKeyWithdrawalRateLimitExemptAmount: "500.00", }) svc := NewWithdrawalService(repo, nil, nil, nil, nil, settingSvc) @@ -71,7 +72,30 @@ func TestWithdrawalServiceSubmitPassesConfiguredRateLimit(t *testing.T) { require.NoError(t, err) require.NotNil(t, repo.submitInput) - require.Equal(t, WithdrawalRateLimitConfig{WindowDays: 7, MaxRequests: 3}, repo.submitInput.RateLimit) + require.Equal(t, WithdrawalRateLimitConfig{WindowDays: 7, MaxRequests: 3, ExemptAmount: 500}, repo.submitInput.RateLimit) +} + +func TestWithdrawalRateLimitExemptsOnlyAmountsStrictlyAboveThreshold(t *testing.T) { + config := WithdrawalRateLimitConfig{WindowDays: 7, MaxRequests: 3, ExemptAmount: 500} + + require.False(t, config.ExemptsAmount(499.99)) + require.False(t, config.ExemptsAmount(500)) + require.True(t, config.ExemptsAmount(500.01)) + + config.ExemptAmount = 0 + require.False(t, config.ExemptsAmount(1000)) +} + +func TestSettingServiceWithdrawalRateLimitRejectsInvalidExemptAmount(t *testing.T) { + svc := newModuleSwitchSettingService(map[string]string{ + SettingKeyWithdrawalRateLimitWindowDays: "7", + SettingKeyWithdrawalRateLimitMax: "3", + SettingKeyWithdrawalRateLimitExemptAmount: "500.001", + }) + + _, err := svc.GetWithdrawalRateLimitConfig(context.Background()) + + require.Equal(t, "WITHDRAWAL_RATE_LIMIT_CONFIG_INVALID", infraerrors.Reason(err)) } func TestWithdrawalServiceListMineExposesOnlyUserSafeRejectionReason(t *testing.T) { diff --git a/backend/internal/service/gateway_multiplatform_test.go b/backend/internal/service/gateway_multiplatform_test.go index c6ca4a7c5..3235d94c9 100644 --- a/backend/internal/service/gateway_multiplatform_test.go +++ b/backend/internal/service/gateway_multiplatform_test.go @@ -2465,8 +2465,8 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) { repo := &mockAccountRepoForPlatform{ accounts: []Account{ - {ID: 1, Platform: PlatformAnthropic, Priority: 1, Status: StatusActive, Schedulable: true, Concurrency: 5}, - {ID: 2, Platform: PlatformAnthropic, Priority: 2, Status: StatusActive, Schedulable: true, Concurrency: 5}, + {ID: 1, Platform: PlatformAnthropic, Type: AccountTypeAPIKey, Priority: 1, Status: StatusActive, Schedulable: true, Concurrency: 5, GroupIDs: []int64{groupID}, Credentials: map[string]any{"api_key": "full-secret"}}, + {ID: 2, Platform: PlatformAnthropic, Type: AccountTypeAPIKey, Priority: 2, Status: StatusActive, Schedulable: true, Concurrency: 5, GroupIDs: []int64{groupID}, Credentials: map[string]any{"api_key": "backup-secret"}}, }, accountsByID: map[int64]*Account{}, } @@ -2501,6 +2501,13 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) { acquireResults: map[int64]bool{1: false}, waitCounts: map[int64]int{1: 0}, } + snapshotCache := &snapshotHydrationCache{ + snapshot: []*Account{ + {ID: 1, Platform: PlatformAnthropic, Type: AccountTypeAPIKey, Priority: 1, Status: StatusActive, Schedulable: true, Concurrency: 5, GroupIDs: []int64{groupID}}, + {ID: 2, Platform: PlatformAnthropic, Type: AccountTypeAPIKey, Priority: 2, Status: StatusActive, Schedulable: true, Concurrency: 5, GroupIDs: []int64{groupID}}, + }, + accounts: repo.accountsByID, + } svc := &GatewayService{ accountRepo: repo, @@ -2508,6 +2515,7 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) { cache: cache, cfg: cfg, concurrencyService: NewConcurrencyService(concurrencyCache), + schedulerSnapshot: NewSchedulerSnapshotService(snapshotCache, nil, repo, groupRepo, nil), } result, err := svc.SelectAccountWithLoadAwareness(ctx, &groupID, sessionHash, "claude-3-5-sonnet-20241022", nil, "", int64(0)) @@ -2515,6 +2523,7 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) { require.NotNil(t, result) require.NotNil(t, result.WaitPlan) require.Equal(t, int64(1), result.Account.ID) + require.Equal(t, "full-secret", result.Account.GetCredential("api_key")) }) t.Run("模型路由-粘性账号命中", func(t *testing.T) { diff --git a/backend/internal/service/gateway_request.go b/backend/internal/service/gateway_request.go index 7d556c1af..bb9766465 100644 --- a/backend/internal/service/gateway_request.go +++ b/backend/internal/service/gateway_request.go @@ -126,6 +126,18 @@ func normalizeSessionUserAgentFallback(raw string) string { return strings.Join(strings.Fields(normalized), " ") } +const claudeCodeLongContextModelSuffix = "[1m]" + +// Claude Code 把 [1m] 当作客户端侧上下文选择器，正常情况下不会发送给上游。 +// 同时兼容客户端重复拼接后缀的情况。 +func normalizeClaudeCodeLongContextModel(model string) string { + for len(model) > len(claudeCodeLongContextModelSuffix) && + strings.EqualFold(model[len(model)-len(claudeCodeLongContextModelSuffix):], claudeCodeLongContextModelSuffix) { + model = model[:len(model)-len(claudeCodeLongContextModelSuffix)] + } + return model +} + // ParseGatewayRequest 解析网关请求体并返回结构化结果。 // protocol 指定请求协议格式（domain.PlatformAnthropic / domain.PlatformGemini）， // 不同协议使用不同的 system/messages 字段名。 @@ -154,6 +166,19 @@ func ParseGatewayRequest(body []byte, protocol string) (*ParsedRequest, error) { return nil, fmt.Errorf("invalid model field type") } parsed.Model = modelResult.String() + if protocol == domain.PlatformAnthropic { + normalizedModel := normalizeClaudeCodeLongContextModel(parsed.Model) + if normalizedModel != parsed.Model { + normalizedBody, err := sjson.SetBytes(body, "model", normalizedModel) + if err != nil { + return nil, fmt.Errorf("normalize model field: %w", err) + } + parsed.Body = normalizedBody + body = normalizedBody + jsonStr = *(*string)(unsafe.Pointer(&body)) + parsed.Model = normalizedModel + } + } } // stream: 需要严格类型校验，非 bool 返回错误 diff --git a/backend/internal/service/gateway_request_test.go b/backend/internal/service/gateway_request_test.go index 40bd11867..0079ad6a4 100644 --- a/backend/internal/service/gateway_request_test.go +++ b/backend/internal/service/gateway_request_test.go @@ -76,6 +76,24 @@ func TestParseGatewayRequest_InvalidStreamType(t *testing.T) { require.Error(t, err) } +func TestParseGatewayRequest_NormalizesClaudeCodeLongContextSuffix(t *testing.T) { + parsed, err := ParseGatewayRequest( + []byte(`{"model":"claude-sonnet-4-5[1m][1M]","messages":[]}`), + domain.PlatformAnthropic, + ) + + require.NoError(t, err) + require.Equal(t, "claude-sonnet-4-5", parsed.Model) + require.JSONEq(t, `{"model":"claude-sonnet-4-5","messages":[]}`, string(parsed.Body)) +} + +func TestParseGatewayRequest_DoesNotNormalizeLongContextSuffixForOtherProtocols(t *testing.T) { + parsed, err := ParseGatewayRequest([]byte(`{"model":"custom[1m]"}`), domain.PlatformOpenAI) + + require.NoError(t, err) + require.Equal(t, "custom[1m]", parsed.Model) +} + // ============ Gemini 原生格式解析测试 ============ func TestParseGatewayRequest_GeminiContents(t *testing.T) { diff --git a/backend/internal/service/gateway_service.go b/backend/internal/service/gateway_service.go index bc0c1bdde..536ea5af0 100644 --- a/backend/internal/service/gateway_service.go +++ b/backend/internal/service/gateway_service.go @@ -139,7 +139,7 @@ func openAIStreamEventIsTerminal(data string) bool { return true } switch gjson.Get(trimmed, "type").String() { - case "response.completed", "response.done", "response.failed", "response.incomplete", "response.cancelled", "response.canceled": + case "response.completed", "response.done", "response.failed", "response.incomplete", "response.cancelled", "response.canceled", "error", "response.error": return true default: return false @@ -489,11 +489,24 @@ type AccountWaitPlan struct { MaxWaiting int } +// OpenAIAccountDispatchRequirements captures the immutable constraints used to +// select an OpenAI-compatible account. Dispatch revalidation must use the same +// effective constraints (including any scheduler fallback) before forwarding. +type OpenAIAccountDispatchRequirements struct { + RequestedModel string + RequiredTransport OpenAIUpstreamTransport + RequiredImageCapability OpenAIImagesCapability + RequiredEndpointCapability OpenAIEndpointCapability + RequiredPlatform string + RequireCompact bool +} + type AccountSelectionResult struct { - Account *Account - Acquired bool - ReleaseFunc func() - WaitPlan *AccountWaitPlan // nil means no wait allowed + Account *Account + Acquired bool + ReleaseFunc func() + WaitPlan *AccountWaitPlan // nil means no wait allowed + OpenAIDispatchRequirements *OpenAIAccountDispatchRequirements } // ClaudeUsage 表示Claude API返回的usage信息 @@ -601,6 +614,16 @@ const ( type GatewayFailureReason string +const ( + // GatewayFailureReasonOpenAIFirstOutputTimeout identifies a native OpenAI + // first-output deadline exceeded after an upstream attempt was started. + GatewayFailureReasonOpenAIFirstOutputTimeout GatewayFailureReason = "openai_first_output_timeout" + // GatewayFailureReasonRoutingBudgetExhausted identifies a local routing + // budget that expired before another upstream attempt could be started. + // It must not poison account health because no account response was observed. + GatewayFailureReasonRoutingBudgetExhausted GatewayFailureReason = "routing_budget_exhausted" +) + // UpstreamFailoverError indicates an upstream or credential error that may // trigger account failover. type UpstreamFailoverError struct { @@ -637,6 +660,9 @@ func (e *UpstreamFailoverError) ShouldReportAccountScheduleFailure() bool { if e == nil { return false } + if e.Reason == GatewayFailureReasonRoutingBudgetExhausted { + return false + } return !e.IsCredentialFailure() || e.Scope == GatewayFailureScopeAccount } @@ -1371,7 +1397,7 @@ func (s *GatewayService) buildOAuthMetadataUserID(parsed *ParsedRequest, account // - account：必须是 OAuth 账号，且调用方已判断不是 Claude Code 客户端。 // - body：已经 marshal 成 Anthropic /v1/messages 格式的请求体。 // - systemRaw：body 中原始 system 字段（用于判断是否需要 rewrite）。 -// - model：最终会发给上游的模型 ID（用于 haiku 旁路 + metadata 版本选择）。 +// - model：最终会发给上游的模型 ID（用于模型规范化 + metadata 版本选择）。 // // 返回：改写后的 body。即使中间任何一步失败，也会退化成原 body（不会 panic）。 func (s *GatewayService) applyClaudeCodeOAuthMimicryToBody( @@ -1388,7 +1414,7 @@ func (s *GatewayService) applyClaudeCodeOAuthMimicryToBody( systemPromptInjectionEnabled, systemPrompt, systemPromptBlocks := s.claudeOAuthSystemPromptInjectionSettings(ctx) systemRewritten := false - if systemPromptInjectionEnabled && !strings.Contains(strings.ToLower(model), "haiku") { + if systemPromptInjectionEnabled { body = rewriteSystemForNonClaudeCodeWithPromptBlocks(body, systemRaw, systemPrompt, systemPromptBlocks) systemRewritten = true } @@ -1861,15 +1887,12 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro stickyCacheMissReason = "session_limit" // 会话限制已满，继续到负载感知选择 } else { - return &AccountSelectionResult{ - Account: stickyAccount, - WaitPlan: &AccountWaitPlan{ - AccountID: stickyAccountID, - MaxConcurrency: stickyAccount.Concurrency, - Timeout: waitTimeoutForStickyDecision(cfg.StickySessionWaitTimeout, stickyRuntime), - MaxWaiting: maxWaiting, - }, - }, nil + return s.newSelectionResult(ctx, stickyAccount, false, nil, &AccountWaitPlan{ + AccountID: stickyAccountID, + MaxConcurrency: stickyAccount.Concurrency, + Timeout: waitTimeoutForStickyDecision(cfg.StickySessionWaitTimeout, stickyRuntime), + MaxWaiting: maxWaiting, + }) } } else { stickyCacheMissReason = "wait_queue_full" @@ -3066,6 +3089,9 @@ func (s *GatewayService) hydrateSelectedAccount(ctx context.Context, account *Ac func (s *GatewayService) newSelectionResult(ctx context.Context, account *Account, acquired bool, release func(), waitPlan *AccountWaitPlan) (*AccountSelectionResult, error) { hydrated, err := s.hydrateSelectedAccount(ctx, account) if err != nil { + if acquired && release != nil { + release() + } return nil, err } return &AccountSelectionResult{ @@ -5162,13 +5188,13 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A // 检测到"有 CC prompt 但无 billing block"的不一致而判为 third-party。 // Parrot 的 transform_request 从不检查客户端 system 内容，直接覆盖。 systemRewritten := false - if systemPromptInjectionEnabled && !strings.Contains(strings.ToLower(reqModel), "haiku") { + if systemPromptInjectionEnabled { body = rewriteSystemForNonClaudeCodeWithPromptBlocks(body, parsed.System, systemPrompt, systemPromptBlocks) systemRewritten = true } // system 被重写时保留 CC prompt 的 cache_control: ephemeral（匹配真实 Claude Code 行为）； - // 未重写时（haiku / 已含 CC 前缀）剥离客户端 cache_control，与原有行为一致。 + // 未重写时（注入开关关闭）剥离客户端 cache_control，与原有行为一致。 // 两种情况下 enforceCacheControlLimit 都会兜底处理上限。 normalizeOpts := claudeOAuthNormalizeOptions{stripSystemCacheControl: !systemRewritten} if s.identityService != nil { @@ -7005,15 +7031,9 @@ func (s *GatewayService) buildUpstreamRequest(ctx context.Context, c *gin.Contex applyClaudeCodeMimicHeaders(req, reqStream) incomingBeta := getHeaderRaw(req.Header, "anthropic-beta") - // Claude Code OAuth credentials are scoped to Claude Code. - // Non-haiku models MUST include claude-code beta for Anthropic to recognize - // this as a legitimate Claude Code request; without it, the request is - // rejected as third-party ("out of extra usage"). - // Haiku models are exempt from third-party detection and don't need it. - requiredBetas := []string{claude.BetaOAuth, claude.BetaInterleavedThinking} - if !strings.Contains(strings.ToLower(modelID), "haiku") { - requiredBetas = claude.FullClaudeCodeMimicryBetas() - } + // OAuth mimic 对所有模型（包括 Haiku）使用完整 Claude Code beta， + // 否则 Haiku 请求仍可能被识别为第三方客户端。 + requiredBetas := claude.FullClaudeCodeMimicryBetas() setHeaderRaw(req.Header, "anthropic-beta", mergeAnthropicBetaDropping(requiredBetas, incomingBeta, effectiveDropSet)) } else { // Claude Code 客户端：尽量透传原始 header，仅补齐 oauth beta diff --git a/backend/internal/service/grok_media.go b/backend/internal/service/grok_media.go index 22850135a..9be544a67 100644 --- a/backend/internal/service/grok_media.go +++ b/backend/internal/service/grok_media.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "encoding/json" + "errors" "fmt" "io" "mime" @@ -30,10 +31,15 @@ const ( GrokMediaEndpointVideosEdits GrokMediaEndpoint = "videos_edits" GrokMediaEndpointVideosExtensions GrokMediaEndpoint = "videos_extensions" GrokMediaEndpointVideoStatus GrokMediaEndpoint = "video_status" + GrokMediaEndpointVideoContent GrokMediaEndpoint = "video_content" ) func (e GrokMediaEndpoint) RequiresRequestBody() bool { - return e != GrokMediaEndpointVideoStatus + return !e.IsVideoLookupRequest() +} + +func (e GrokMediaEndpoint) IsVideoLookupRequest() bool { + return e == GrokMediaEndpointVideoStatus || e == GrokMediaEndpointVideoContent } func (e GrokMediaEndpoint) IsGenerationRequest() bool { @@ -113,7 +119,7 @@ func (r GrokMediaRequestInfo) ModerationBody() []byte { } func (e GrokMediaEndpoint) httpMethod() string { - if e == GrokMediaEndpointVideoStatus { + if e.IsVideoLookupRequest() { return http.MethodGet } return http.MethodPost @@ -272,12 +278,13 @@ func parseGrokMediaMultipartRequest(contentType string, body []byte, info *GrokM } } -func GrokMediaVideoRequestSessionHash(requestID string) string { +func GrokMediaVideoRequestSessionHash(requestID string, userID, apiKeyID int64) string { requestID = strings.TrimSpace(requestID) - if requestID == "" { + if requestID == "" || userID <= 0 || apiKeyID <= 0 { return "" } - return "grok-video:" + DeriveSessionHashFromSeed(requestID) + ownerSeed := fmt.Sprintf("%d:%d:%s", userID, apiKeyID, requestID) + return "grok-video:" + DeriveSessionHashFromSeed(ownerSeed) } func GrokMediaSessionHash(sessionHash string) string { @@ -288,8 +295,88 @@ func GrokMediaSessionHash(sessionHash string) string { return "grok-media:" + sessionHash } -func (s *OpenAIGatewayService) BindGrokMediaVideoRequestAccount(ctx context.Context, groupID *int64, requestID string, accountID int64) error { - return s.BindStickySession(ctx, groupID, GrokMediaVideoRequestSessionHash(requestID), accountID) +const grokMediaVideoOwnerBindingPrefix = "grok-video-owner:" + +func grokMediaVideoOwnerBindingKey(requestID string, userID, apiKeyID int64) string { + sessionHash := GrokMediaVideoRequestSessionHash(requestID, userID, apiKeyID) + if sessionHash == "" { + return "" + } + return grokMediaVideoOwnerBindingPrefix + sessionHash +} + +func (s *OpenAIGatewayService) grokMediaVideoBindingTTL() time.Duration { + if s != nil && s.cfg != nil && s.cfg.Gateway.OpenAIWS.StickySessionTTLSeconds > 0 { + return time.Duration(s.cfg.Gateway.OpenAIWS.StickySessionTTLSeconds) * time.Second + } + return openaiStickySessionTTL +} + +func (s *OpenAIGatewayService) BindGrokMediaVideoRequestAccount( + ctx context.Context, + groupID *int64, + requestID string, + userID, apiKeyID, accountID int64, +) error { + if s == nil || s.cache == nil { + return fmt.Errorf("grok video request binding cache is unavailable") + } + ownerKey := grokMediaVideoOwnerBindingKey(requestID, userID, apiKeyID) + routingKey := s.openAISessionCacheKey(GrokMediaVideoRequestSessionHash(requestID, userID, apiKeyID)) + if ownerKey == "" || routingKey == "" || accountID <= 0 { + return fmt.Errorf("grok video request binding is invalid") + } + ttl := s.grokMediaVideoBindingTTL() + group := derefGroupID(groupID) + if err := s.cache.SetSessionString(ctx, group, ownerKey, strconv.FormatInt(accountID, 10), ttl); err != nil { + return fmt.Errorf("store grok video owner binding: %w", err) + } + if err := s.cache.SetSessionAccountID(ctx, group, routingKey, accountID, ttl); err != nil { + return fmt.Errorf("store grok video routing binding: %w", err) + } + return nil +} + +func (s *OpenAIGatewayService) ResolveGrokMediaVideoRequestAccount( + ctx context.Context, + groupID *int64, + requestID string, + userID, apiKeyID int64, +) (int64, error) { + if s == nil || s.cache == nil { + return 0, fmt.Errorf("grok video request binding cache is unavailable") + } + ownerKey := grokMediaVideoOwnerBindingKey(requestID, userID, apiKeyID) + routingKey := s.openAISessionCacheKey(GrokMediaVideoRequestSessionHash(requestID, userID, apiKeyID)) + if ownerKey == "" || routingKey == "" { + return 0, fmt.Errorf("grok video request binding is invalid") + } + group := derefGroupID(groupID) + ttl := s.grokMediaVideoBindingTTL() + rawAccountID, err := s.cache.GetSessionString(ctx, group, ownerKey) + if err != nil && !errors.Is(err, ErrGatewaySessionStringNotFound) { + return 0, fmt.Errorf("load grok video owner binding: %w", err) + } + if errors.Is(err, ErrGatewaySessionStringNotFound) { + // Compatibility for in-flight tasks created before owner bindings used + // their own namespace. Migrate the old routing record on first lookup. + legacyAccountID, legacyErr := s.cache.GetSessionAccountID(ctx, group, routingKey) + if legacyErr != nil { + return 0, legacyErr + } + rawAccountID = strconv.FormatInt(legacyAccountID, 10) + } + accountID, parseErr := strconv.ParseInt(strings.TrimSpace(rawAccountID), 10, 64) + if parseErr != nil || accountID <= 0 { + return 0, fmt.Errorf("grok video owner binding is invalid") + } + if err := s.cache.SetSessionString(ctx, group, ownerKey, strconv.FormatInt(accountID, 10), ttl); err != nil { + return 0, fmt.Errorf("refresh grok video owner binding: %w", err) + } + if err := s.cache.SetSessionAccountID(ctx, group, routingKey, accountID, ttl); err != nil { + return 0, fmt.Errorf("restore grok video routing binding: %w", err) + } + return accountID, nil } func (s *OpenAIGatewayService) ForwardGrokMedia( @@ -313,6 +400,9 @@ func (s *OpenAIGatewayService) ForwardGrokMedia( if err != nil { return nil, err } + if endpoint == GrokMediaEndpointVideoContent { + return s.forwardGrokMediaVideoContent(ctx, c, account, token, requestID, startTime) + } targetURL, err := buildGrokMediaURL(ctx, account, s.cfg, s.settingService, endpoint, requestID) if err != nil { return nil, err @@ -385,6 +475,13 @@ func (s *OpenAIGatewayService) ForwardGrokMedia( if err != nil { return nil, err } + if endpoint == GrokMediaEndpointVideoStatus { + respBody = rewriteGrokMediaVideoContentURLs( + respBody, + requestID, + grokMediaContentProxyURL(c, requestID), + ) + } writeGrokMediaResponse(c, resp, respBody, s.responseHeaderFilter) usage := grokMediaUsageFromResponse(endpoint, requestInfo, respBody) return &OpenAIForwardResult{ @@ -406,6 +503,156 @@ func (s *OpenAIGatewayService) ForwardGrokMedia( }, nil } +func (s *OpenAIGatewayService) forwardGrokMediaVideoContent( + ctx context.Context, + c *gin.Context, + account *Account, + token, requestID string, + startTime time.Time, +) (*OpenAIForwardResult, error) { + statusURL, err := buildGrokMediaURL( + ctx, account, s.cfg, s.settingService, GrokMediaEndpointVideoStatus, requestID, + ) + if err != nil { + return nil, err + } + + upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx) + defer releaseUpstreamCtx() + statusReq, err := http.NewRequestWithContext( + WithHTTPUpstreamRedirectsDisabled(upstreamCtx), + http.MethodGet, + statusURL, + nil, + ) + if err != nil { + return nil, err + } + statusReq.Header.Set("Authorization", "Bearer "+token) + statusReq.Header.Set("Accept", "application/json") + if account.IsGrokOAuth() && isGrokCLIProxyTarget(statusURL) { + applyGrokCLIHeaders(statusReq.Header) + } + account.ApplyHeaderOverrides(statusReq.Header) + + proxyURL := "" + if account.ProxyID != nil && account.Proxy != nil { + proxyURL = account.Proxy.URL() + } + upstreamStart := time.Now() + statusResp, err := s.httpUpstream.Do(statusReq, proxyURL, account.ID, account.Concurrency) + if err != nil { + SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds()) + return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, false) + } + statusRequestID := firstNonEmpty(statusResp.Header.Get("x-request-id"), statusResp.Header.Get("xai-request-id")) + if statusResp.StatusCode >= http.StatusMultipleChoices { + defer func() { _ = statusResp.Body.Close() }() + SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds()) + if statusResp.StatusCode < http.StatusBadRequest { + return nil, fmt.Errorf("grok media status redirect is not allowed") + } + return s.handleGrokMediaErrorResponse(ctx, statusResp, c, account, statusRequestID, "") + } + statusBody, err := ReadUpstreamResponseBody(statusResp.Body, s.cfg, c, openAITooLargeError) + _ = statusResp.Body.Close() + if err != nil { + SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds()) + return nil, err + } + s.updateGrokUsageSnapshot(ctx, account.ID, xai.ParseQuotaHeaders(statusResp.Header, statusResp.StatusCode)) + + contentURL, err := grokMediaSignedVideoContentURL(statusBody, requestID) + if err != nil { + SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds()) + return nil, err + } + signedContent := contentURL != "" + if !signedContent { + contentURL, err = buildGrokMediaURL( + ctx, account, s.cfg, s.settingService, GrokMediaEndpointVideoContent, requestID, + ) + if err != nil { + SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds()) + return nil, err + } + } + + contentReq, err := http.NewRequestWithContext( + WithHTTPUpstreamRedirectsDisabled(upstreamCtx), + http.MethodGet, + contentURL, + nil, + ) + if err != nil { + SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds()) + return nil, err + } + contentReq.Header.Set("Accept", "*/*") + if c != nil { + if rangeHeader := strings.TrimSpace(c.GetHeader("Range")); rangeHeader != "" { + contentReq.Header.Set("Range", rangeHeader) + } + } + if !signedContent { + contentReq.Header.Set("Authorization", "Bearer "+token) + if account.IsGrokOAuth() && isGrokCLIProxyTarget(contentURL) { + applyGrokCLIHeaders(contentReq.Header) + } + account.ApplyHeaderOverrides(contentReq.Header) + } + + contentResp, err := s.httpUpstream.Do(contentReq, proxyURL, account.ID, account.Concurrency) + SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds()) + if err != nil { + return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, false) + } + defer func() { _ = contentResp.Body.Close() }() + contentRequestID := firstNonEmpty( + contentResp.Header.Get("x-request-id"), + contentResp.Header.Get("xai-request-id"), + statusRequestID, + ) + if contentResp.StatusCode >= http.StatusMultipleChoices && contentResp.StatusCode < http.StatusBadRequest { + return nil, fmt.Errorf("grok media content redirect is not allowed") + } + if contentResp.StatusCode >= http.StatusBadRequest && contentResp.StatusCode != http.StatusRequestedRangeNotSatisfiable { + return s.handleGrokMediaErrorResponse(ctx, contentResp, c, account, contentRequestID, "") + } + + s.updateGrokUsageSnapshot(ctx, account.ID, xai.ParseQuotaHeaders(contentResp.Header, contentResp.StatusCode)) + if err := writeGrokMediaContentResponse(c, contentResp); err != nil { + return nil, err + } + return &OpenAIForwardResult{ + RequestID: contentRequestID, + ResponseHeaders: contentResp.Header.Clone(), + Duration: time.Since(startTime), + }, nil +} + +func grokMediaSignedVideoContentURL(body []byte, requestID string) (string, error) { + rawURL := strings.TrimSpace(gjson.GetBytes(body, "video.url").String()) + if rawURL == "" { + return "", nil + } + if isGrokMediaVideoContentURL(rawURL, requestID) { + return "", nil + } + parsed, err := url.Parse(rawURL) + if err != nil || !strings.EqualFold(parsed.Scheme, "https") || + !strings.EqualFold(parsed.Hostname(), "vidgen.x.ai") || + (parsed.Port() != "" && parsed.Port() != "443") || parsed.User != nil { + return "", fmt.Errorf("grok media status returned an unsupported video content URL") + } + return parsed.String(), nil +} + +func isGrokCLIProxyTarget(rawURL string) bool { + parsed, err := url.Parse(strings.TrimSpace(rawURL)) + return err == nil && strings.EqualFold(parsed.Hostname(), "cli-chat-proxy.grok.com") +} + func prepareGrokMediaForwardBody(endpoint GrokMediaEndpoint, body []byte, contentType string) ([]byte, string, error) { if endpoint != GrokMediaEndpointImagesEdits || gjson.ValidBytes(body) { return body, contentType, nil @@ -722,3 +969,140 @@ func writeGrokMediaResponse(c *gin.Context, resp *http.Response, body []byte, fi } c.Data(resp.StatusCode, contentType, body) } + +func writeGrokMediaContentResponse(c *gin.Context, resp *http.Response) error { + if c == nil || resp == nil || resp.Body == nil { + return fmt.Errorf("grok media content response is incomplete") + } + + for _, name := range []string{ + "Content-Type", + "Content-Length", + "Content-Range", + "Accept-Ranges", + "Content-Disposition", + } { + if value := strings.TrimSpace(resp.Header.Get(name)); value != "" { + c.Header(name, value) + } + } + if strings.TrimSpace(c.Writer.Header().Get("Content-Length")) == "" && resp.ContentLength >= 0 { + c.Header("Content-Length", strconv.FormatInt(resp.ContentLength, 10)) + } + if strings.TrimSpace(c.Writer.Header().Get("Content-Type")) == "" { + c.Header("Content-Type", "application/octet-stream") + } + c.Status(resp.StatusCode) + MarkResponseCommitted(c) + _, err := io.Copy(c.Writer, resp.Body) + return err +} + +func rewriteGrokMediaVideoContentURLs(body []byte, requestID, proxyURL string) []byte { + if len(body) == 0 || strings.TrimSpace(requestID) == "" || strings.TrimSpace(proxyURL) == "" || !gjson.ValidBytes(body) { + return body + } + + decoder := json.NewDecoder(bytes.NewReader(body)) + decoder.UseNumber() + var value any + if err := decoder.Decode(&value); err != nil { + return body + } + changed := rewriteGrokMediaKnownVideoURL(&value, proxyURL) + if rewriteGrokMediaVideoContentURLValue(&value, requestID, proxyURL) { + changed = true + } + if !changed { + return body + } + rewritten, err := json.Marshal(value) + if err != nil { + return body + } + return rewritten +} + +func rewriteGrokMediaKnownVideoURL(value *any, proxyURL string) bool { + if value == nil { + return false + } + root, ok := (*value).(map[string]any) + if !ok { + return false + } + video, ok := root["video"].(map[string]any) + if !ok { + return false + } + rawURL, ok := video["url"].(string) + if !ok || strings.TrimSpace(rawURL) == "" { + return false + } + video["url"] = proxyURL + return true +} + +func rewriteGrokMediaVideoContentURLValue(value *any, requestID, proxyURL string) bool { + if value == nil { + return false + } + switch typed := (*value).(type) { + case map[string]any: + changed := false + for key, child := range typed { + childValue := child + if rewriteGrokMediaVideoContentURLValue(&childValue, requestID, proxyURL) { + typed[key] = childValue + changed = true + } + } + return changed + case []any: + changed := false + for index, child := range typed { + childValue := child + if rewriteGrokMediaVideoContentURLValue(&childValue, requestID, proxyURL) { + typed[index] = childValue + changed = true + } + } + return changed + case string: + if isGrokMediaVideoContentURL(typed, requestID) { + *value = proxyURL + return true + } + } + return false +} + +func isGrokMediaVideoContentURL(rawURL, requestID string) bool { + parsed, err := url.Parse(strings.TrimSpace(rawURL)) + if err != nil || parsed.Path == "" { + return false + } + segments := strings.Split(strings.Trim(parsed.EscapedPath(), "/"), "/") + if len(segments) < 3 { + return false + } + requestID = strings.Trim(requestID, "/") + decodedID, err := url.PathUnescape(segments[len(segments)-2]) + if err != nil { + return false + } + return segments[len(segments)-3] == "videos" && + decodedID == requestID && + segments[len(segments)-1] == "content" +} + +func grokMediaContentProxyURL(c *gin.Context, requestID string) string { + if c == nil || c.Request == nil || c.Request.URL == nil || strings.TrimSpace(requestID) == "" { + return "" + } + pathPrefix := "" + if strings.HasPrefix(c.Request.URL.Path, "/v1/") { + pathPrefix = "/v1" + } + return pathPrefix + "/videos/" + url.PathEscape(strings.Trim(requestID, "/")) + "/content" +} diff --git a/backend/internal/service/grok_media_content_test.go b/backend/internal/service/grok_media_content_test.go new file mode 100644 index 000000000..2f2d2bc5a --- /dev/null +++ b/backend/internal/service/grok_media_content_test.go @@ -0,0 +1,294 @@ +package service + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +type grokMediaContentUpstreamStub struct { + requests []*http.Request + responses []*http.Response +} + +func (s *grokMediaContentUpstreamStub) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) { + s.requests = append(s.requests, req) + if len(s.responses) == 0 { + return nil, io.EOF + } + resp := s.responses[0] + s.responses = s.responses[1:] + return resp, nil +} + +func (s *grokMediaContentUpstreamStub) DoWithTLS( + req *http.Request, + proxyURL string, + accountID int64, + accountConcurrency int, + _ *tlsfingerprint.Profile, +) (*http.Response, error) { + return s.Do(req, proxyURL, accountID, accountConcurrency) +} + +func grokMediaContentTestAccount() *Account { + return &Account{ + ID: 9, + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Credentials: map[string]any{ + "api_key": "upstream-key", + "base_url": "https://relay.example/v1", + }, + } +} + +func grokMediaContentTestContext(method, target string, headers map[string]string) (*gin.Context, *httptest.ResponseRecorder) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(method, target, nil) + for name, value := range headers { + c.Request.Header.Set(name, value) + } + return c, recorder +} + +func grokMediaContentStatusResponse(body string) *http.Response { + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(body)), + } +} + +func TestForwardGrokMediaContentFetchesSignedURLWithoutCredentials(t *testing.T) { + upstream := &grokMediaContentUpstreamStub{ + responses: []*http.Response{ + grokMediaContentStatusResponse(`{"status":"done","video":{"url":"https://vidgen.x.ai/signed-token/task-1.mp4?signature=secret"}}`), + { + StatusCode: http.StatusPartialContent, + Header: http.Header{ + "Content-Type": []string{"video/mp4"}, + "Content-Length": []string{"13"}, + "Content-Range": []string{"bytes 0-12/100"}, + "Accept-Ranges": []string{"bytes"}, + "Content-Disposition": []string{`attachment; filename="task-1.mp4"`}, + }, + Body: io.NopCloser(strings.NewReader("video-payload")), + }, + }, + } + svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream} + c, recorder := grokMediaContentTestContext( + http.MethodGet, + "https://api.example/v1/videos/task-1/content", + map[string]string{"Range": "bytes=0-12"}, + ) + + result, err := svc.ForwardGrokMedia( + context.Background(), c, grokMediaContentTestAccount(), + GrokMediaEndpointVideoContent, "task-1", nil, "", + ) + + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, http.StatusPartialContent, recorder.Code) + require.Equal(t, "video-payload", recorder.Body.String()) + require.Len(t, upstream.requests, 2) + require.Equal(t, "https://relay.example/v1/videos/task-1", upstream.requests[0].URL.String()) + require.Equal(t, "Bearer upstream-key", upstream.requests[0].Header.Get("Authorization")) + require.True(t, HTTPUpstreamRedirectsDisabled(upstream.requests[0].Context())) + require.Equal(t, "https://vidgen.x.ai/signed-token/task-1.mp4?signature=secret", upstream.requests[1].URL.String()) + require.Empty(t, upstream.requests[1].Header.Get("Authorization")) + require.Equal(t, "bytes=0-12", upstream.requests[1].Header.Get("Range")) + require.True(t, HTTPUpstreamRedirectsDisabled(upstream.requests[1].Context())) + require.Equal(t, "video/mp4", recorder.Header().Get("Content-Type")) + require.Equal(t, "bytes 0-12/100", recorder.Header().Get("Content-Range")) + require.Equal(t, "bytes", recorder.Header().Get("Accept-Ranges")) + require.Equal(t, `attachment; filename="task-1.mp4"`, recorder.Header().Get("Content-Disposition")) + require.True(t, IsResponseCommitted(c)) +} + +func TestForwardGrokMediaContentFollowsAuthenticatedSub2APIChain(t *testing.T) { + for _, statusURL := range []string{ + `/v1/videos/task-1/content`, + `https://different-relay.example/v1/videos/task-1/content`, + } { + t.Run(statusURL, func(t *testing.T) { + upstream := &grokMediaContentUpstreamStub{ + responses: []*http.Response{ + grokMediaContentStatusResponse(`{"status":"completed","video":{"url":"` + statusURL + `"}}`), + { + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"video/mp4"}}, + Body: io.NopCloser(strings.NewReader("video-payload")), + }, + }, + } + svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream} + c, recorder := grokMediaContentTestContext( + http.MethodGet, "https://api.example/v1/videos/task-1/content", nil, + ) + + _, err := svc.ForwardGrokMedia( + context.Background(), c, grokMediaContentTestAccount(), + GrokMediaEndpointVideoContent, "task-1", nil, "", + ) + + require.NoError(t, err) + require.Equal(t, http.StatusOK, recorder.Code) + require.Equal(t, "video-payload", recorder.Body.String()) + require.Len(t, upstream.requests, 2) + require.Equal(t, "https://relay.example/v1/videos/task-1/content", upstream.requests[1].URL.String()) + require.Equal(t, "Bearer upstream-key", upstream.requests[1].Header.Get("Authorization")) + }) + } +} + +func TestForwardGrokMediaContentPreservesRangeNotSatisfiable(t *testing.T) { + upstream := &grokMediaContentUpstreamStub{ + responses: []*http.Response{ + grokMediaContentStatusResponse(`{"status":"completed"}`), + { + StatusCode: http.StatusRequestedRangeNotSatisfiable, + Header: http.Header{ + "Content-Type": []string{"text/plain"}, + "Content-Range": []string{"bytes */100"}, + "Accept-Ranges": []string{"bytes"}, + }, + Body: io.NopCloser(strings.NewReader("bad-range")), + }, + }, + } + svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream} + c, recorder := grokMediaContentTestContext( + http.MethodGet, + "https://api.example/v1/videos/task-1/content", + map[string]string{"Range": "bytes=500-600"}, + ) + + _, err := svc.ForwardGrokMedia( + context.Background(), c, grokMediaContentTestAccount(), + GrokMediaEndpointVideoContent, "task-1", nil, "", + ) + + require.NoError(t, err) + require.Equal(t, http.StatusRequestedRangeNotSatisfiable, recorder.Code) + require.Equal(t, "bad-range", recorder.Body.String()) + require.Equal(t, "bytes */100", recorder.Header().Get("Content-Range")) + require.Equal(t, "bytes", recorder.Header().Get("Accept-Ranges")) +} + +func TestForwardGrokMediaContentRejectsRedirectResponses(t *testing.T) { + for _, tt := range []struct { + name string + responses []*http.Response + wantCalls int + }{ + { + name: "status redirect", + responses: []*http.Response{{ + StatusCode: http.StatusFound, + Header: http.Header{"Location": []string{"https://attacker.invalid/status"}}, + Body: http.NoBody, + }}, + wantCalls: 1, + }, + { + name: "content redirect", + responses: []*http.Response{ + grokMediaContentStatusResponse(`{"status":"done","video":{"url":"https://vidgen.x.ai/task-1.mp4"}}`), + { + StatusCode: http.StatusTemporaryRedirect, + Header: http.Header{"Location": []string{"https://attacker.invalid/content"}}, + Body: http.NoBody, + }, + }, + wantCalls: 2, + }, + } { + t.Run(tt.name, func(t *testing.T) { + upstream := &grokMediaContentUpstreamStub{responses: tt.responses} + svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream} + c, _ := grokMediaContentTestContext( + http.MethodGet, "https://api.example/v1/videos/task-1/content", nil, + ) + + _, err := svc.ForwardGrokMedia( + context.Background(), c, grokMediaContentTestAccount(), + GrokMediaEndpointVideoContent, "task-1", nil, "", + ) + + require.ErrorContains(t, err, "redirect is not allowed") + require.Len(t, upstream.requests, tt.wantCalls) + }) + } +} + +func TestGrokMediaSignedVideoContentURLValidation(t *testing.T) { + valid, err := grokMediaSignedVideoContentURL( + []byte(`{"video":{"url":"https://vidgen.x.ai/video.mp4?signature=secret"}}`), + "task-1", + ) + require.NoError(t, err) + require.Equal(t, "https://vidgen.x.ai/video.mp4?signature=secret", valid) + + relay, err := grokMediaSignedVideoContentURL( + []byte(`{"video":{"url":"/v1/videos/task-1/content"}}`), + "task-1", + ) + require.NoError(t, err) + require.Empty(t, relay) + + for _, rawURL := range []string{ + "http://vidgen.x.ai/video.mp4", + "https://vidgen.x.ai.attacker.invalid/video.mp4", + "https://vidgen.x.ai@attacker.invalid/video.mp4", + "https://vidgen.x.ai:444/video.mp4", + "/v1/videos/task-2/content", + } { + t.Run(rawURL, func(t *testing.T) { + _, err := grokMediaSignedVideoContentURL( + []byte(`{"video":{"url":"`+rawURL+`"}}`), + "task-1", + ) + require.ErrorContains(t, err, "unsupported video content URL") + }) + } +} + +func TestForwardGrokVideoStatusRewritesProtectedURLsToSameOrigin(t *testing.T) { + upstream := &grokMediaContentUpstreamStub{ + responses: []*http.Response{ + grokMediaContentStatusResponse(`{"id":"task/one","status":"completed","video":{"url":"https://vidgen.x.ai/signed.mp4"},"nested":[{"url":"https://relay.example/v1/videos/task%2Fone/content"},{"url":"https://relay.example/v1/videos/other/content"}],"counter":9007199254740993}`), + }, + } + svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream} + c, recorder := grokMediaContentTestContext( + http.MethodGet, + "https://api.example/v1/videos/task%2Fone", + map[string]string{"X-Forwarded-Host": "malicious.invalid"}, + ) + + _, err := svc.ForwardGrokMedia( + context.Background(), c, grokMediaContentTestAccount(), + GrokMediaEndpointVideoStatus, "task/one", nil, "", + ) + + require.NoError(t, err) + require.Equal(t, "/v1/videos/task%2Fone/content", gjson.GetBytes(recorder.Body.Bytes(), "video.url").String()) + require.Equal(t, "/v1/videos/task%2Fone/content", gjson.GetBytes(recorder.Body.Bytes(), "nested.0.url").String()) + require.Equal(t, "https://relay.example/v1/videos/other/content", gjson.GetBytes(recorder.Body.Bytes(), "nested.1.url").String()) + require.Equal(t, "9007199254740993", gjson.GetBytes(recorder.Body.Bytes(), "counter").String()) + require.NotContains(t, recorder.Body.String(), "malicious.invalid") +} diff --git a/backend/internal/service/grok_media_test.go b/backend/internal/service/grok_media_test.go index 1271a6b38..e585a0da4 100644 --- a/backend/internal/service/grok_media_test.go +++ b/backend/internal/service/grok_media_test.go @@ -3,6 +3,7 @@ package service import ( "bytes" "context" + "fmt" "mime/multipart" "net/http" "net/http/httptest" @@ -29,6 +30,172 @@ func TestGrokVideoMutationEndpointsAreBillableGenerationRequests(t *testing.T) { } require.False(t, GrokMediaEndpointVideoStatus.IsGenerationRequest()) require.False(t, GrokMediaEndpointVideoStatus.IsVideoMutationRequest()) + require.True(t, GrokMediaEndpointVideoStatus.IsVideoLookupRequest()) + require.True(t, GrokMediaEndpointVideoContent.IsVideoLookupRequest()) + require.False(t, GrokMediaEndpointVideoContent.RequiresRequestBody()) +} + +func TestGrokMediaVideoRequestSessionHashIsOwnerScoped(t *testing.T) { + t.Parallel() + + base := GrokMediaVideoRequestSessionHash("request-1", 10, 20) + require.NotEmpty(t, base) + require.NotEqual(t, base, GrokMediaVideoRequestSessionHash("request-1", 11, 20)) + require.NotEqual(t, base, GrokMediaVideoRequestSessionHash("request-1", 10, 21)) + require.NotEqual(t, base, GrokMediaVideoRequestSessionHash("request-2", 10, 20)) + require.Empty(t, GrokMediaVideoRequestSessionHash("", 10, 20)) + require.Empty(t, GrokMediaVideoRequestSessionHash("request-1", 0, 20)) + require.Empty(t, GrokMediaVideoRequestSessionHash("request-1", 10, 0)) +} + +type grokOwnerBindingCache struct { + GatewayCache + bindings map[string]int64 + strings map[string]string + ttl time.Duration +} + +func (c *grokOwnerBindingCache) key(groupID int64, sessionHash string) string { + return fmt.Sprintf("%d:%s", groupID, sessionHash) +} + +func (c *grokOwnerBindingCache) GetSessionAccountID(_ context.Context, groupID int64, sessionHash string) (int64, error) { + if accountID, ok := c.bindings[c.key(groupID, sessionHash)]; ok { + return accountID, nil + } + return 0, ErrGatewaySessionStringNotFound +} + +func (c *grokOwnerBindingCache) SetSessionAccountID(_ context.Context, groupID int64, sessionHash string, accountID int64, ttl time.Duration) error { + if c.bindings == nil { + c.bindings = make(map[string]int64) + } + c.bindings[c.key(groupID, sessionHash)] = accountID + c.ttl = ttl + return nil +} + +func (c *grokOwnerBindingCache) DeleteSessionAccountID(_ context.Context, groupID int64, sessionHash string) error { + delete(c.bindings, c.key(groupID, sessionHash)) + return nil +} + +func (c *grokOwnerBindingCache) GetSessionString(_ context.Context, groupID int64, sessionHash string) (string, error) { + if value, ok := c.strings[c.key(groupID, sessionHash)]; ok { + return value, nil + } + return "", ErrGatewaySessionStringNotFound +} + +func (c *grokOwnerBindingCache) SetSessionString(_ context.Context, groupID int64, sessionHash, value string, ttl time.Duration) error { + if c.strings == nil { + c.strings = make(map[string]string) + } + c.strings[c.key(groupID, sessionHash)] = value + c.ttl = ttl + return nil +} + +func TestGrokMediaVideoRequestBindingRejectsOtherOwnersAndGroups(t *testing.T) { + t.Parallel() + + cache := &grokOwnerBindingCache{} + cfg := &config.Config{} + cfg.Gateway.OpenAIWS.StickySessionTTLSeconds = 90 + svc := &OpenAIGatewayService{cache: cache, cfg: cfg} + groupID := int64(7) + require.NoError(t, svc.BindGrokMediaVideoRequestAccount( + context.Background(), &groupID, "request-1", 10, 20, 30, + )) + require.Equal(t, 90*time.Second, cache.ttl) + + accountID, err := svc.ResolveGrokMediaVideoRequestAccount( + context.Background(), &groupID, "request-1", 10, 20, + ) + require.NoError(t, err) + require.Equal(t, int64(30), accountID) + + otherGroupID := int64(8) + for _, lookup := range []struct { + groupID *int64 + userID int64 + apiKeyID int64 + }{ + {groupID: &groupID, userID: 11, apiKeyID: 20}, + {groupID: &groupID, userID: 10, apiKeyID: 21}, + {groupID: &otherGroupID, userID: 10, apiKeyID: 20}, + } { + accountID, err = svc.ResolveGrokMediaVideoRequestAccount( + context.Background(), lookup.groupID, "request-1", lookup.userID, lookup.apiKeyID, + ) + require.Error(t, err) + require.Zero(t, accountID) + } + + _, err = (&OpenAIGatewayService{}).ResolveGrokMediaVideoRequestAccount( + context.Background(), &groupID, "request-1", 10, 20, + ) + require.ErrorContains(t, err, "cache is unavailable") +} + +func TestGrokMediaVideoOwnerBindingSurvivesRoutingEviction(t *testing.T) { + t.Parallel() + + cache := &grokOwnerBindingCache{} + svc := &OpenAIGatewayService{cache: cache} + groupID := int64(7) + const ( + requestID = "request-recover" + userID = int64(10) + apiKeyID = int64(20) + accountID = int64(30) + ) + require.NoError(t, svc.BindGrokMediaVideoRequestAccount( + context.Background(), &groupID, requestID, userID, apiKeyID, accountID, + )) + + ownerKey := grokMediaVideoOwnerBindingKey(requestID, userID, apiKeyID) + routingKey := svc.openAISessionCacheKey(GrokMediaVideoRequestSessionHash(requestID, userID, apiKeyID)) + require.NotEqual(t, ownerKey, routingKey) + require.Equal(t, "30", cache.strings[cache.key(groupID, ownerKey)]) + require.Equal(t, accountID, cache.bindings[cache.key(groupID, routingKey)]) + + // Scheduler eviction removes only the routing/sticky record when the + // account is temporarily blocked. Ownership must remain authoritative. + require.NoError(t, cache.DeleteSessionAccountID(context.Background(), groupID, routingKey)) + _, routingExists := cache.bindings[cache.key(groupID, routingKey)] + require.False(t, routingExists) + + resolvedID, err := svc.ResolveGrokMediaVideoRequestAccount( + context.Background(), &groupID, requestID, userID, apiKeyID, + ) + require.NoError(t, err) + require.Equal(t, accountID, resolvedID) + require.Equal(t, accountID, cache.bindings[cache.key(groupID, routingKey)], "lookup should restore scheduler routing after cooldown") +} + +func TestGrokMediaVideoOwnerBindingMigratesLegacyRoutingRecord(t *testing.T) { + t.Parallel() + + cache := &grokOwnerBindingCache{} + svc := &OpenAIGatewayService{cache: cache} + groupID := int64(7) + const ( + requestID = "request-legacy" + userID = int64(10) + apiKeyID = int64(20) + accountID = int64(30) + ) + routingKey := svc.openAISessionCacheKey(GrokMediaVideoRequestSessionHash(requestID, userID, apiKeyID)) + require.NoError(t, cache.SetSessionAccountID(context.Background(), groupID, routingKey, accountID, time.Minute)) + + resolvedID, err := svc.ResolveGrokMediaVideoRequestAccount( + context.Background(), &groupID, requestID, userID, apiKeyID, + ) + require.NoError(t, err) + require.Equal(t, accountID, resolvedID) + ownerKey := grokMediaVideoOwnerBindingKey(requestID, userID, apiKeyID) + require.Equal(t, "30", cache.strings[cache.key(groupID, ownerKey)]) } func TestParseGrokMediaJSONVideoBillingMetadata(t *testing.T) { @@ -132,6 +299,10 @@ func TestBuildGrokMediaURLSupportsVideoMutationsAndKeepsOAuthOnOfficialAPI(t *te extensionURL, err := buildGrokMediaURL(context.Background(), account, nil, nil, GrokMediaEndpointVideosExtensions, "") require.NoError(t, err) require.Equal(t, "https://api.x.ai/v1/videos/extensions", extensionURL) + + contentURL, err := buildGrokMediaURL(context.Background(), account, nil, nil, GrokMediaEndpointVideoContent, "task/one") + require.NoError(t, err) + require.Equal(t, "https://api.x.ai/v1/videos/task%2Fone/content", contentURL) } type grokURLSettingRepo struct { diff --git a/backend/internal/service/grok_quota_service.go b/backend/internal/service/grok_quota_service.go index 8d5e7ece5..3bd20e616 100644 --- a/backend/internal/service/grok_quota_service.go +++ b/backend/internal/service/grok_quota_service.go @@ -226,6 +226,23 @@ func (s *GrokQuotaService) ProbeBilling(ctx context.Context, accountID int64) (* }) } +// ProbeMediaEligibility refreshes the billing observation used by media +// scheduling and then evaluates the persisted account state. Transport errors +// remain fail-closed; deterministic states such as forbidden or Free are +// returned as normal ineligibility decisions. +func (s *GrokQuotaService) ProbeMediaEligibility(ctx context.Context, accountID int64) (bool, string, error) { + _, probeErr := s.ProbeBilling(ctx, accountID) + account, err := s.loadGrokOAuthAccount(ctx, accountID) + if err != nil { + return false, "billing_probe_failed", err + } + eligible, reason := account.GrokMediaGenerationEligibility() + if reason == "billing_unobserved" && probeErr != nil { + return false, reason, probeErr + } + return eligible, reason, nil +} + func (s *GrokQuotaService) probeBilling(ctx context.Context, accountID int64) (*GrokQuotaProbeResult, error) { account, token, proxyURL, err := s.prepareProbe(ctx, accountID) if err != nil { @@ -254,12 +271,25 @@ func (s *GrokQuotaService) probeBilling(ctx context.Context, accountID int64) (* weeklyOK := weekly.summary != nil monthlyOK := monthly.summary != nil + previous, _ := grokBillingSnapshotFromExtra(account.Extra) if !weeklyOK && !monthlyOK { - return nil, mergeGrokBillingProbeErrors(weekly.status, monthly.status, weekly.err, monthly.err) + probeErr := mergeGrokBillingProbeErrors(weekly.status, monthly.status, weekly.err, monthly.err) + billing := xai.MergeBillingProbeResult(previous, nil, nil, false, false) + if billing == nil { + billing = &xai.BillingSummary{Partial: true, FailedWindows: []string{"weekly", "monthly"}} + } + billing.WeeklyStatusCode = weekly.status + billing.MonthlyStatusCode = monthly.status + billing = xai.StampBillingSummary(billing, preferBillingObservationStatus(weekly.status, monthly.status), "billing_probe") + if persistErr := s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{grokBillingExtraKey: billing}); persistErr != nil { + slog.Warn("grok_billing_failure_persist_failed", "account_id", account.ID, "error", persistErr) + } + return nil, probeErr } statusCode := preferSuccessfulBillingStatus(weekly.status, monthly.status, weeklyOK, monthlyOK) - previous, _ := grokBillingSnapshotFromExtra(account.Extra) billing := xai.MergeBillingProbeResult(previous, weekly.summary, monthly.summary, weeklyOK, monthlyOK) + billing.WeeklyStatusCode = weekly.status + billing.MonthlyStatusCode = monthly.status billing = xai.StampBillingSummary(billing, statusCode, "billing_probe") persistErr := s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{ grokBillingExtraKey: billing, @@ -281,6 +311,16 @@ func (s *GrokQuotaService) probeBilling(ctx context.Context, accountID int64) (* }, nil } +func preferBillingObservationStatus(weeklyStatus, monthlyStatus int) int { + if weeklyStatus == http.StatusForbidden || monthlyStatus == http.StatusForbidden { + return http.StatusForbidden + } + if weeklyStatus != 0 { + return weeklyStatus + } + return monthlyStatus +} + func (s *GrokQuotaService) runProbeFlight( ctx context.Context, key string, diff --git a/backend/internal/service/grok_upstream_url.go b/backend/internal/service/grok_upstream_url.go index f30d673fd..ce1f980c5 100644 --- a/backend/internal/service/grok_upstream_url.go +++ b/backend/internal/service/grok_upstream_url.go @@ -102,6 +102,12 @@ func buildGrokMediaURL( return xai.BuildVideosExtensionsURLWithValidator(baseURL, validator) case GrokMediaEndpointVideoStatus: return xai.BuildVideoURLWithValidator(baseURL, requestID, validator) + case GrokMediaEndpointVideoContent: + videoURL, err := xai.BuildVideoURLWithValidator(baseURL, requestID, validator) + if err != nil { + return "", err + } + return videoURL + "/content", nil default: return "", fmt.Errorf("unsupported grok media endpoint: %s", endpoint) } diff --git a/backend/internal/service/http_upstream_profile.go b/backend/internal/service/http_upstream_profile.go index 2d63bbd5e..c9e128799 100644 --- a/backend/internal/service/http_upstream_profile.go +++ b/backend/internal/service/http_upstream_profile.go @@ -12,6 +12,7 @@ const ( ) type httpUpstreamProfileContextKey struct{} +type httpUpstreamDisableRedirectsContextKey struct{} // WithHTTPUpstreamProfile injects an upstream transport profile into ctx. func WithHTTPUpstreamProfile(ctx context.Context, profile HTTPUpstreamProfile) context.Context { @@ -40,3 +41,18 @@ func HTTPUpstreamProfileFromContext(ctx context.Context) HTTPUpstreamProfile { return HTTPUpstreamProfileDefault } } + +// WithHTTPUpstreamRedirectsDisabled prevents credential-bearing requests from +// following redirects through a shared upstream client. +func WithHTTPUpstreamRedirectsDisabled(ctx context.Context) context.Context { + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, httpUpstreamDisableRedirectsContextKey{}, true) +} + +// HTTPUpstreamRedirectsDisabled reports whether redirects must be returned to +// the caller instead of being followed by the upstream HTTP client. +func HTTPUpstreamRedirectsDisabled(ctx context.Context) bool { + return ctx != nil && ctx.Value(httpUpstreamDisableRedirectsContextKey{}) == true +} diff --git a/backend/internal/service/http_upstream_profile_test.go b/backend/internal/service/http_upstream_profile_test.go index 284bedcd2..eeb8bee6f 100644 --- a/backend/internal/service/http_upstream_profile_test.go +++ b/backend/internal/service/http_upstream_profile_test.go @@ -25,3 +25,13 @@ func TestHTTPUpstreamProfileRejectsUnknownValue(t *testing.T) { t.Fatalf("profile = %q, want default", profile) } } + +func TestWithHTTPUpstreamRedirectsDisabled(t *testing.T) { + ctx := WithHTTPUpstreamRedirectsDisabled(nil) + if !HTTPUpstreamRedirectsDisabled(ctx) { + t.Fatal("redirects should be disabled") + } + if HTTPUpstreamRedirectsDisabled(context.Background()) { + t.Fatal("redirects should remain enabled by default") + } +} diff --git a/backend/internal/service/image_generation_intent.go b/backend/internal/service/image_generation_intent.go index 8dd7637d7..439015488 100644 --- a/backend/internal/service/image_generation_intent.go +++ b/backend/internal/service/image_generation_intent.go @@ -60,6 +60,26 @@ func IsImageGenerationIntent(endpoint string, requestedModel string, body []byte return openAIJSONToolChoiceSelectsImageGeneration(gjson.GetBytes(body, "tool_choice")) } +// IsExplicitImageGenerationIntent only classifies request signals that +// explicitly require image generation. Passive Codex image_gen namespace or +// function declarations remain available to the model, but must not trigger +// permission, Responses-capability, or image-concurrency gates by themselves. +func IsExplicitImageGenerationIntent(endpoint string, requestedModel string, body []byte) bool { + if IsImageGenerationEndpoint(endpoint) || isOpenAIImageGenerationModel(requestedModel) { + return true + } + if len(body) == 0 || !gjson.ValidBytes(body) { + return false + } + if model := strings.TrimSpace(gjson.GetBytes(body, "model").String()); isOpenAIImageGenerationModel(model) { + return true + } + if openAIJSONToolsContainNativeImageGeneration(gjson.GetBytes(body, "tools")) { + return true + } + return openAIJSONToolChoiceSelectsExplicitImageGeneration(gjson.GetBytes(body, "tool_choice")) +} + // 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) { @@ -117,6 +137,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(item.Get("type").String()) + return !found + }) + return found +} + func isOpenAIImageGenerationType(value string) bool { return strings.TrimSpace(value) == "image_generation" } @@ -175,6 +207,39 @@ 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(choice.Get("namespace").String(), choice.Get("name").String()) { + return true + } + if fn := choice.Get("function"); fn.IsObject() { + return isOpenAIImageGenFunctionReference(fn.Get("namespace").String(), fn.Get("name").String()) + } + 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_explicit_test.go b/backend/internal/service/image_generation_intent_explicit_test.go new file mode 100644 index 000000000..9827601af --- /dev/null +++ b/backend/internal/service/image_generation_intent_explicit_test.go @@ -0,0 +1,94 @@ +package service + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestIsExplicitImageGenerationIntent(t *testing.T) { + tests := []struct { + name string + endpoint string + requestedModel string + body string + want bool + }{ + { + name: "passive namespace declaration", + body: `{"model":"gpt-5.5","tools":[{"type":"namespace","name":"image_gen","tools":[{"type":"function","name":"imagegen"}]}],"tool_choice":"auto","input":"write code"}`, + }, + { + name: "passive Responses Lite declaration", + body: `{"model":"gpt-5.5","input":[{"type":"additional_tools","tools":[{"type":"namespace","name":"image_gen","tools":[{"type":"function","name":"imagegen"}]}]},{"type":"message","role":"user","content":"write code"}],"tool_choice":"auto"}`, + }, + { + name: "passive flattened function declaration", + body: `{"model":"gpt-5.5","tools":[{"type":"function","name":"image_gen.imagegen"}],"tool_choice":"auto","input":"write code"}`, + }, + { + name: "native image tool declaration", + body: `{"model":"gpt-5.5","tools":[{"type":"image_generation","model":"gpt-image-2"}],"tool_choice":"auto"}`, + want: true, + }, + { + name: "requested image model", + requestedModel: "gpt-image-2", + want: true, + }, + { + name: "image endpoint", + endpoint: "/v1/images/generations", + want: true, + }, + { + name: "namespace tool choice", + body: `{"model":"gpt-5.5","tools":[{"type":"namespace","name":"image_gen"}],"tool_choice":{"type":"namespace","name":"image_gen"}}`, + want: true, + }, + { + name: "flattened function tool choice", + body: `{"model":"gpt-5.5","tools":[{"type":"function","name":"image_gen.imagegen"}],"tool_choice":{"type":"function","name":"image_gen.imagegen"}}`, + want: true, + }, + { + name: "wrapped function tool choice", + body: `{"model":"gpt-5.5","tool_choice":{"tool":{"type":"function","name":"image_gen__imagegen"}}}`, + want: true, + }, + { + name: "function object tool choice", + body: `{"model":"gpt-5.5","tool_choice":{"type":"function","function":{"namespace":"image_gen","name":"imagegen"}}}`, + want: true, + }, + { + name: "image call history is not current intent", + body: `{"model":"gpt-5.5","input":[{"type":"function_call","namespace":"image_gen","name":"imagegen","arguments":"{}"},{"type":"image_generation_call","id":"ig_1"}]}`, + }, + { + name: "malformed body", + body: `{"model":`, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + endpoint := test.endpoint + if endpoint == "" { + endpoint = openAIResponsesEndpoint + } + require.Equal(t, test.want, IsExplicitImageGenerationIntent( + endpoint, + test.requestedModel, + []byte(test.body), + )) + }) + } +} + +func TestIsExplicitImageGenerationIntentKeepsGeneralDeclarationDetectionSeparate(t *testing.T) { + body := []byte(`{"model":"gpt-5.5","tools":[{"type":"namespace","name":"image_gen"}],"tool_choice":"auto"}`) + + require.True(t, IsImageGenerationIntent(openAIResponsesEndpoint, "gpt-5.5", body)) + require.False(t, IsExplicitImageGenerationIntent(openAIResponsesEndpoint, "gpt-5.5", body)) +} 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 000000000..8a03a31ed --- /dev/null +++ b/backend/internal/service/openai_account_model_transient.go @@ -0,0 +1,164 @@ +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) 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 000000000..5a56e17a6 --- /dev/null +++ b/backend/internal/service/openai_account_model_transient_test.go @@ -0,0 +1,72 @@ +package service + +import ( + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestOpenAIAccountModelTransientStateIsolatesModelsAndAccounts(t *testing.T) { + now := time.Date(2026, 7, 22, 0, 0, 0, 0, time.UTC) + state := newOpenAIAccountModelTransientState(8) + + first := state.recordFailure(7, "gpt-5.6", now) + require.Equal(t, 1, first.FailureStreak) + require.Zero(t, first.Cooldown) + require.False(t, state.isBlocked(7, "gpt-5.6", now)) + + second := state.recordFailure(7, "gpt-5.6", now.Add(time.Second)) + require.Equal(t, 2, second.FailureStreak) + require.Equal(t, openAIModelTransientShortCooldown, second.Cooldown) + require.True(t, state.isBlocked(7, "gpt-5.6", now.Add(2*time.Second))) + require.False(t, state.isBlocked(7, "gpt-5.5", now.Add(2*time.Second))) + require.False(t, state.isBlocked(8, "gpt-5.6", now.Add(2*time.Second))) + require.False(t, state.isBlocked(7, "gpt-5.6", second.BlockUntil.Add(time.Nanosecond))) +} + +func TestOpenAIGatewayModelTransientUsesMappedModel(t *testing.T) { + svc := &OpenAIGatewayService{openaiModelTransient: newOpenAIAccountModelTransientState(8)} + account := &Account{ + ID: 11, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Credentials: map[string]any{ + "model_mapping": map[string]any{"gpt-alias": "gpt-5.6"}, + }, + } + now := time.Now() + + svc.recordOpenAIAccountModelTransientFailure(account, "gpt-alias", now) + svc.recordOpenAIAccountModelTransientFailure(account, "gpt-5.6", now.Add(time.Second)) + + require.True(t, svc.isOpenAIAccountModelRuntimeBlocked(account, "gpt-alias")) + require.True(t, svc.isOpenAIAccountModelRuntimeBlocked(account, "gpt-5.6")) + require.False(t, svc.isOpenAIAccountModelRuntimeBlocked(account, "gpt-5.5")) +} + +func TestReportOpenAIAccountScheduleSuccessClearsCanonicalModelWithoutRemapping(t *testing.T) { + svc := &OpenAIGatewayService{openaiModelTransient: newOpenAIAccountModelTransientState(8)} + account := &Account{ + ID: 12, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Credentials: map[string]any{ + "model_mapping": map[string]any{ + "gpt-alias": "gpt-mapped", + "gpt-mapped": "gpt-remapped", + }, + }, + } + now := time.Now() + + svc.recordOpenAIAccountModelTransientFailure(account, "gpt-alias", now) + svc.recordOpenAIAccountModelTransientFailure(account, "gpt-alias", now.Add(time.Second)) + require.True(t, svc.isOpenAIAccountModelRuntimeBlocked(account, "gpt-alias")) + + canonicalModel := account.GetMappedModel("gpt-alias") + require.Equal(t, "gpt-mapped", canonicalModel) + svc.ReportOpenAIAccountScheduleResult(account.ID, true, nil, canonicalModel) + + require.False(t, svc.isOpenAIAccountModelRuntimeBlocked(account, "gpt-alias")) +} diff --git a/backend/internal/service/openai_account_runtime_block.go b/backend/internal/service/openai_account_runtime_block.go index 5487d4e71..55e0026a9 100644 --- a/backend/internal/service/openai_account_runtime_block.go +++ b/backend/internal/service/openai_account_runtime_block.go @@ -2,6 +2,9 @@ package service import ( "context" + "log/slog" + "net/http" + "strings" "time" ) @@ -60,3 +63,102 @@ func (s *OpenAIGatewayService) isOpenAIAccountRuntimeBlocked(account *Account) b s.openaiAccountRuntimeBlockUntil.Delete(account.ID) 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 (s *OpenAIGatewayService) recordOpenAIAccountModelTransientFailure(account *Account, requestedModel string, now time.Time) openAIAccountModelTransientDecision { + if s == nil || account == nil { + return openAIAccountModelTransientDecision{} + } + state := s.getOpenAIAccountModelTransientState() + if state == nil { + return openAIAccountModelTransientDecision{} + } + model := canonicalOpenAIAccountSchedulingModel(account, requestedModel) + return state.recordFailure(account.ID, model, now) +} + +func (s *OpenAIGatewayService) isOpenAIAccountModelRuntimeBlocked(account *Account, requestedModel string) bool { + if s == nil || account == nil { + return false + } + state := s.getOpenAIAccountModelTransientState() + if state == nil { + return false + } + model := canonicalOpenAIAccountSchedulingModel(account, requestedModel) + return state.isBlocked(account.ID, model, time.Now()) +} + +func (s *OpenAIGatewayService) isOpenAIAccountRequestRuntimeBlocked(account *Account, requestedModel string) bool { + return s != nil && (s.isOpenAIAccountRuntimeBlocked(account) || s.isOpenAIAccountModelRuntimeBlocked(account, requestedModel)) +} + +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) handleOpenAIAccountUpstreamErrorForModel( + ctx context.Context, + account *Account, + requestedModel string, + statusCode int, + headers http.Header, + responseBody []byte, +) bool { + if s == nil || account == nil { + return false + } + stateCtx, cancel := openAIAccountStateContext(ctx) + defer cancel() + + if account.Platform == PlatformOpenAI && account.Type == AccountTypeAPIKey && strings.TrimSpace(requestedModel) != "" && shouldCooldownOpenAITransientUpstreamError(statusCode, responseBody) { + decision := s.recordOpenAIAccountModelTransientFailure(account, requestedModel, time.Now()) + if decision.FailureStreak > 0 { + slog.Warn("openai_model_transient_state", + "account_id", account.ID, + "model", canonicalOpenAIAccountSchedulingModel(account, requestedModel), + "failure_streak", decision.FailureStreak, + "cooldown_ms", decision.Cooldown.Milliseconds(), + "block_scope", "account_model", + ) + } + return false + } + if s.rateLimitService == nil { + return false + } + shouldDisable := s.rateLimitService.HandleUpstreamErrorForModel(stateCtx, account, requestedModel, statusCode, headers, responseBody) + if shouldDisable { + s.BlockAccountScheduling(account, time.Time{}, "upstream_disable") + } + return shouldDisable +} diff --git a/backend/internal/service/openai_account_scheduler.go b/backend/internal/service/openai_account_scheduler.go index 763194167..88b81b86f 100644 --- a/backend/internal/service/openai_account_scheduler.go +++ b/backend/internal/service/openai_account_scheduler.go @@ -52,15 +52,16 @@ var openAIAdvancedSchedulerSettingCache atomic.Value // *cachedOpenAIAdvancedSch var openAIAdvancedSchedulerSettingSF singleflight.Group type OpenAIAccountScheduleRequest struct { - GroupID *int64 - SessionHash string - StickyAccountID int64 - PreviousResponseID string - RequestedModel string - RequiredTransport OpenAIUpstreamTransport - RequiredImageCapability OpenAIImagesCapability - RequireCompact bool - ExcludedIDs map[int64]struct{} + GroupID *int64 + SessionHash string + StickyAccountID int64 + PreviousResponseID string + RequestedModel string + RequiredTransport OpenAIUpstreamTransport + RequiredImageCapability OpenAIImagesCapability + RequiredEndpointCapability OpenAIEndpointCapability + RequireCompact bool + ExcludedIDs map[int64]struct{} } type OpenAIAccountScheduleDecision struct { @@ -279,7 +280,8 @@ func (s *defaultOpenAIAccountScheduler) Select( return nil, decision, err } if selection != nil && selection.Account != nil { - if !s.isAccountTransportCompatible(selection.Account, req.RequiredTransport) { + if !s.isAccountTransportCompatible(selection.Account, req.RequiredTransport) || + !s.isAccountRequestCompatible(selection.Account, req) { if selection.ReleaseFunc != nil { selection.ReleaseFunc() } @@ -368,7 +370,8 @@ func (s *defaultOpenAIAccountScheduler) selectBySessionHash( return nil, nil } account = s.service.recheckSelectedOpenAIAccountFromDB(ctx, req.GroupID, account, req.RequestedModel, req.RequireCompact) - if account == nil || !s.isAccountTransportCompatible(account, req.RequiredTransport) { + if account == nil || !s.isAccountTransportCompatible(account, req.RequiredTransport) || + !s.isAccountRequestCompatible(account, req) { _ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, sessionHash) return nil, nil } @@ -382,6 +385,13 @@ func (s *defaultOpenAIAccountScheduler) selectBySessionHash( ReleaseFunc: result.ReleaseFunc, }, nil } + // 普通 HTTP/Chat 请求没有 previous_response_id，切换到同组的健康账号 + // 不会破坏上游会话连续性。不要让一个繁忙的粘性账号把请求阻塞到 + // StickySessionWaitTimeout（生产默认 120 秒）；继续走负载均衡层， + // 由候选账号立即尝试可用槽位。只有 WS continuation 才必须保留等待计划。 + if strings.TrimSpace(req.PreviousResponseID) == "" { + return nil, nil + } cfg := s.service.schedulingConfig() // WaitPlan.MaxConcurrency 使用 Concurrency（非 EffectiveLoadFactor），因为 WaitPlan 控制的是 Redis 实际并发槽位等待。 @@ -1139,6 +1149,9 @@ func (s *defaultOpenAIAccountScheduler) filterOpenAIAccountsForLoadBalance( if !account.IsSchedulable() || !account.IsOpenAICompatible() { continue } + if s.service != nil && s.service.isOpenAIAccountRequestRuntimeBlocked(account, req.RequestedModel) { + continue + } // require_privacy_set: 跳过 privacy 未设置的账号并标记异常 if schedGroup != nil && schedGroup.RequirePrivacySet && !account.IsPrivacySet() { _ = s.service.accountRepo.SetError(ctx, account.ID, @@ -1174,10 +1187,14 @@ func (s *defaultOpenAIAccountScheduler) isAccountRequestCompatible(account *Acco if account == nil { return false } + if s != nil && s.service != nil && s.service.isOpenAIAccountRequestRuntimeBlocked(account, req.RequestedModel) { + return false + } if req.RequestedModel != "" && !account.IsModelSupported(req.RequestedModel) { return false } - return accountSupportsRequestedOpenAIImageCapability(account, req.RequiredImageCapability) + return accountSupportsRequestedOpenAIImageCapability(account, req.RequiredImageCapability) && + account.SupportsOpenAIEndpointCapability(req.RequiredEndpointCapability) } func (s *defaultOpenAIAccountScheduler) ReportResult(accountID int64, success bool, firstTokenMs *int) { @@ -1300,7 +1317,7 @@ func (s *OpenAIGatewayService) SelectAccountWithScheduler( requiredTransport OpenAIUpstreamTransport, requireCompact bool, ) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error) { - return s.selectAccountWithScheduler(ctx, groupID, previousResponseID, sessionHash, requestedModel, excludedIDs, requiredTransport, "", requireCompact) + return s.selectAccountWithScheduler(ctx, groupID, previousResponseID, sessionHash, requestedModel, excludedIDs, requiredTransport, "", "", requireCompact) } func (s *OpenAIGatewayService) SelectAccountWithSchedulerForImages( @@ -1311,17 +1328,38 @@ func (s *OpenAIGatewayService) SelectAccountWithSchedulerForImages( excludedIDs map[int64]struct{}, requiredCapability OpenAIImagesCapability, ) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error) { - selection, decision, err := s.selectAccountWithScheduler(ctx, groupID, "", sessionHash, requestedModel, excludedIDs, OpenAIUpstreamTransportHTTPSSE, requiredCapability, false) + selection, decision, err := s.selectAccountWithScheduler(ctx, groupID, "", sessionHash, requestedModel, excludedIDs, OpenAIUpstreamTransportHTTPSSE, requiredCapability, "", false) if err == nil && selection != nil && selection.Account != nil { + setOpenAIImagesDispatchRequirements(selection, requestedModel, requiredCapability) return selection, decision, nil } // 如果要求 native 能力（如指定了模型）但没有可用的 APIKey 账号，回退到 basic（OAuth 账号） if requiredCapability == OpenAIImagesCapabilityNative { - return s.selectAccountWithScheduler(ctx, groupID, "", sessionHash, requestedModel, excludedIDs, OpenAIUpstreamTransportHTTPSSE, OpenAIImagesCapabilityBasic, false) + selection, decision, err = s.selectAccountWithScheduler(ctx, groupID, "", sessionHash, requestedModel, excludedIDs, OpenAIUpstreamTransportHTTPSSE, OpenAIImagesCapabilityBasic, "", false) + if err == nil && selection != nil && selection.Account != nil { + setOpenAIImagesDispatchRequirements(selection, requestedModel, OpenAIImagesCapabilityBasic) + } + return selection, decision, err } return selection, decision, err } +func setOpenAIImagesDispatchRequirements(selection *AccountSelectionResult, requestedModel string, capability OpenAIImagesCapability) { + if selection == nil || selection.Account == nil { + return + } + if capability == OpenAIImagesCapabilityNative && + !selection.Account.SupportsOpenAIImageCapability(OpenAIImagesCapabilityNative) && + selection.Account.SupportsOpenAIImageCapability(OpenAIImagesCapabilityBasic) { + capability = OpenAIImagesCapabilityBasic + } + selection.OpenAIDispatchRequirements = &OpenAIAccountDispatchRequirements{ + RequestedModel: requestedModel, + RequiredTransport: OpenAIUpstreamTransportHTTPSSE, + RequiredImageCapability: capability, + } +} + func (s *OpenAIGatewayService) selectAccountShareModeBoundAccount( ctx context.Context, groupID *int64, @@ -1329,6 +1367,7 @@ func (s *OpenAIGatewayService) selectAccountShareModeBoundAccount( excludedIDs map[int64]struct{}, requiredTransport OpenAIUpstreamTransport, requiredImageCapability OpenAIImagesCapability, + requiredEndpointCapability OpenAIEndpointCapability, requireCompact bool, ) (*AccountSelectionResult, OpenAIAccountScheduleDecision, bool, error) { decision := OpenAIAccountScheduleDecision{Layer: openAIAccountScheduleLayerAccountShareMode} @@ -1403,7 +1442,7 @@ func (s *OpenAIGatewayService) selectAccountShareModeBoundAccount( if !retryCurrentMembership && requestedModel != "" && !account.IsModelSupported(requestedModel) { return nil, decision, true, accountShareModeUnsupportedModelError(requestedModel) } - if !retryCurrentMembership && !isOpenAIAccountEligibleForRequest(account, requestedModel, requireCompact) { + if !retryCurrentMembership && (!isOpenAIAccountEligibleForRequest(account, requestedModel, requireCompact) || s.isOpenAIAccountRequestRuntimeBlocked(account, requestedModel)) { lastErr = ErrNoAvailableAccounts retryCurrentMembership = true } @@ -1411,6 +1450,10 @@ func (s *OpenAIGatewayService) selectAccountShareModeBoundAccount( lastErr = ErrNoAvailableAccounts retryCurrentMembership = true } + if !retryCurrentMembership && !account.SupportsOpenAIEndpointCapability(requiredEndpointCapability) { + lastErr = ErrNoAvailableAccounts + retryCurrentMembership = true + } if !retryCurrentMembership && !s.isOpenAIAccountTransportCompatible(account, requiredTransport) { lastErr = ErrNoAvailableAccounts retryCurrentMembership = true @@ -1486,10 +1529,11 @@ func (s *OpenAIGatewayService) selectAccountWithScheduler( excludedIDs map[int64]struct{}, requiredTransport OpenAIUpstreamTransport, requiredImageCapability OpenAIImagesCapability, + requiredEndpointCapability OpenAIEndpointCapability, requireCompact bool, ) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error) { decision := OpenAIAccountScheduleDecision{} - if selection, accountModeDecision, handled, err := s.selectAccountShareModeBoundAccount(ctx, groupID, requestedModel, excludedIDs, requiredTransport, requiredImageCapability, requireCompact); handled { + if selection, accountModeDecision, handled, err := s.selectAccountShareModeBoundAccount(ctx, groupID, requestedModel, excludedIDs, requiredTransport, requiredImageCapability, requiredEndpointCapability, requireCompact); handled { return selection, accountModeDecision, err } scheduler := s.getOpenAIAccountScheduler(ctx) @@ -1498,14 +1542,15 @@ func (s *OpenAIGatewayService) selectAccountWithScheduler( if requiredTransport == OpenAIUpstreamTransportAny || requiredTransport == OpenAIUpstreamTransportHTTPSSE { effectiveExcludedIDs := cloneExcludedAccountIDs(excludedIDs) for { - selection, err := s.selectAccountWithLoadAwareness(ctx, groupID, sessionHash, requestedModel, effectiveExcludedIDs, requireCompact) + selection, err := s.selectAccountWithLoadAwarenessForRequest(ctx, groupID, sessionHash, requestedModel, effectiveExcludedIDs, requireCompact, previousResponseID != "") if err != nil { return nil, decision, err } if selection == nil || selection.Account == nil { return selection, decision, nil } - if accountSupportsRequestedOpenAIImageCapability(selection.Account, requiredImageCapability) { + if accountSupportsRequestedOpenAIImageCapability(selection.Account, requiredImageCapability) && + selection.Account.SupportsOpenAIEndpointCapability(requiredEndpointCapability) { return selection, decision, nil } if selection.ReleaseFunc != nil { @@ -1523,14 +1568,15 @@ func (s *OpenAIGatewayService) selectAccountWithScheduler( effectiveExcludedIDs := cloneExcludedAccountIDs(excludedIDs) for { - selection, err := s.selectAccountWithLoadAwareness(ctx, groupID, sessionHash, requestedModel, effectiveExcludedIDs, requireCompact) + selection, err := s.selectAccountWithLoadAwarenessForRequest(ctx, groupID, sessionHash, requestedModel, effectiveExcludedIDs, requireCompact, previousResponseID != "") if err != nil { return nil, decision, err } if selection == nil || selection.Account == nil { return selection, decision, nil } - if s.isOpenAIAccountTransportCompatible(selection.Account, requiredTransport) { + if s.isOpenAIAccountTransportCompatible(selection.Account, requiredTransport) && + selection.Account.SupportsOpenAIEndpointCapability(requiredEndpointCapability) { return selection, decision, nil } if selection.ReleaseFunc != nil { @@ -1554,15 +1600,16 @@ func (s *OpenAIGatewayService) selectAccountWithScheduler( } return scheduler.Select(ctx, OpenAIAccountScheduleRequest{ - GroupID: groupID, - SessionHash: sessionHash, - StickyAccountID: stickyAccountID, - PreviousResponseID: previousResponseID, - RequestedModel: requestedModel, - RequiredTransport: requiredTransport, - RequiredImageCapability: requiredImageCapability, - RequireCompact: requireCompact, - ExcludedIDs: excludedIDs, + GroupID: groupID, + SessionHash: sessionHash, + StickyAccountID: stickyAccountID, + PreviousResponseID: previousResponseID, + RequestedModel: requestedModel, + RequiredTransport: requiredTransport, + RequiredImageCapability: requiredImageCapability, + RequiredEndpointCapability: requiredEndpointCapability, + RequireCompact: requireCompact, + ExcludedIDs: excludedIDs, }) } @@ -1587,7 +1634,12 @@ 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, success bool, firstTokenMs *int, canonicalModels ...string) { + if success && len(canonicalModels) > 0 { + if state := s.getOpenAIAccountModelTransientState(); state != nil { + state.recordSuccess(accountID, canonicalModels[0]) + } + } 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 65b15d2c8..7b81b333f 100644 --- a/backend/internal/service/openai_account_scheduler_test.go +++ b/backend/internal/service/openai_account_scheduler_test.go @@ -367,6 +367,103 @@ func TestOpenAIGatewayService_SelectAccountWithScheduler_DefaultDisabledUsesLega require.False(t, decision.StickyPreviousHit) } +func TestOpenAIGatewayService_SelectAccountWithSchedulerForImages_FallbackPersistsBasicDispatchRequirement(t *testing.T) { + resetOpenAIAdvancedSchedulerSettingCacheForTest() + + ctx := context.Background() + groupID := int64(10107) + account := Account{ + ID: 36003, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + GroupIDs: []int64{groupID}, + } + cfg := &config.Config{} + cfg.Gateway.Scheduling.LoadBatchEnabled = false + svc := &OpenAIGatewayService{ + accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{account}}, + cache: &schedulerTestGatewayCache{}, + cfg: cfg, + concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}), + } + + selection, _, err := svc.SelectAccountWithSchedulerForImages( + ctx, + &groupID, + "", + "gpt-image-2", + nil, + OpenAIImagesCapabilityNative, + ) + require.NoError(t, err) + require.NotNil(t, selection) + require.NotNil(t, selection.Account) + require.Equal(t, account.ID, selection.Account.ID) + if selection.ReleaseFunc != nil { + defer selection.ReleaseFunc() + } + + require.NotNil(t, selection.OpenAIDispatchRequirements) + require.Equal(t, "gpt-image-2", selection.OpenAIDispatchRequirements.RequestedModel) + require.Equal(t, OpenAIUpstreamTransportHTTPSSE, selection.OpenAIDispatchRequirements.RequiredTransport) + require.Equal(t, OpenAIImagesCapabilityBasic, selection.OpenAIDispatchRequirements.RequiredImageCapability) + + latest, err := svc.RevalidateSelectedOpenAIAccountForDispatch(ctx, &groupID, selection.Account, *selection.OpenAIDispatchRequirements) + require.NoError(t, err) + require.NotNil(t, latest) + require.Equal(t, account.ID, latest.ID) +} + +func TestOpenAIGatewayService_SelectAccountWithSchedulerForImages_AccountShareModeOAuthPersistsBasicDispatchRequirement(t *testing.T) { + modeGroupID := int64(10110) + consumerUserID := int64(5581) + apiKeyID := int64(20104) + account := Account{ + ID: 36004, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + } + shareRepo := &accountShareModeRepoStub{ + membership: &AccountShareMembership{ID: 2, AccountID: account.ID, ConsumerUserID: consumerUserID, APIKeyID: apiKeyID}, + listing: &AccountShareListing{ID: 2, AccountID: account.ID, OwnerUserID: 1, Status: AccountShareListingStatusActive}, + } + svc := &OpenAIGatewayService{ + accountRepo: stubOpenAIAccountRepo{accounts: []Account{account}}, + accountShareModeService: &AccountShareModeService{repo: shareRepo}, + } + ctx := WithAccountShareModeRequest(context.Background(), consumerUserID, apiKeyID) + + selection, decision, err := svc.SelectAccountWithSchedulerForImages( + ctx, + &modeGroupID, + "", + "gpt-image-2", + nil, + OpenAIImagesCapabilityNative, + ) + require.NoError(t, err) + require.NotNil(t, selection) + require.NotNil(t, selection.Account) + require.Equal(t, openAIAccountScheduleLayerAccountShareMode, decision.Layer) + if selection.ReleaseFunc != nil { + defer selection.ReleaseFunc() + } + + require.NotNil(t, selection.OpenAIDispatchRequirements) + require.Equal(t, OpenAIImagesCapabilityBasic, selection.OpenAIDispatchRequirements.RequiredImageCapability) + + latest, err := svc.RevalidateSelectedOpenAIAccountForDispatch(ctx, &modeGroupID, selection.Account, *selection.OpenAIDispatchRequirements) + require.NoError(t, err) + require.NotNil(t, latest) + require.Equal(t, account.ID, latest.ID) +} + func TestOpenAIGatewayService_SelectAccountWithScheduler_DefaultDisabled_RequiredWSV2_SkipsHTTPOnlyAccount(t *testing.T) { resetOpenAIAdvancedSchedulerSettingCacheForTest() @@ -834,7 +931,7 @@ func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionSticky(t *testin } } -func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionStickyBusyEscapesWhenQueueFull(t *testing.T) { +func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionStickyBusyEscapesForReplaySafeRequest(t *testing.T) { ctx := context.Background() groupID := int64(10100) accounts := []Account{ @@ -876,7 +973,7 @@ func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionStickyBusyEscape 21002: true, }, waitCounts: map[int64]int{ - 21001: 999, + 21001: 0, // 普通请求不应因粘性账号繁忙而进入长等待 }, loadMap: map[int64]*AccountLoadInfo{ 21001: {AccountID: 21001, LoadRate: 90, WaitingCount: 9}, diff --git a/backend/internal/service/openai_alpha_search.go b/backend/internal/service/openai_alpha_search.go index 5aef983d8..ccb27dca0 100644 --- a/backend/internal/service/openai_alpha_search.go +++ b/backend/internal/service/openai_alpha_search.go @@ -100,12 +100,14 @@ func (s *OpenAIGatewayService) ForwardAlphaSearch(ctx context.Context, c *gin.Co upstreamMessage := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(respBody))) if s.shouldFailoverOpenAIUpstreamResponse(resp.StatusCode, upstreamMessage, respBody) { resp.Body = io.NopCloser(bytes.NewReader(respBody)) - s.handleFailoverSideEffects(ctx, resp, account) - return nil, &UpstreamFailoverError{ - StatusCode: resp.StatusCode, - ResponseBody: respBody, - RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), - } + s.handleFailoverSideEffectsForModel(ctx, resp, account, requestedModel) + return nil, newOpenAIUpstreamFailoverError( + resp.StatusCode, + resp.Header, + respBody, + upstreamMessage, + account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), + ) } } diff --git a/backend/internal/service/openai_codex_models_service.go b/backend/internal/service/openai_codex_models_service.go index 5d382abcd..bc661ca06 100644 --- a/backend/internal/service/openai_codex_models_service.go +++ b/backend/internal/service/openai_codex_models_service.go @@ -1,8 +1,10 @@ package service import ( + "bytes" "context" "crypto/sha256" + "encoding/json" "errors" "fmt" "io" @@ -44,6 +46,7 @@ type codexModelsManifestUpstreamError struct { err error retryable bool statusCode int + headers http.Header body []byte } @@ -306,6 +309,7 @@ func (s *OpenAIGatewayService) FetchCodexModelsManifest(ctx context.Context, acc manifest, fetchErr := s.fetchCodexModelsManifestUpstream(ctx, request, ifNoneMatch) if !account.IsOpenAIAgentIdentity() || !isAgentIdentityTaskInvalidCodexModelsError(fetchErr) { + s.handleCodexModelsManifestAccountAuthError(ctx, account, request, fetchErr) return manifest, fetchErr } expectedTaskID := strings.TrimSpace(account.GetCredential("task_id")) @@ -334,6 +338,21 @@ func isAgentIdentityTaskInvalidCodexModelsError(err error) bool { isAgentIdentityTaskInvalidHTTPResponse(upstreamErr.statusCode, upstreamErr.body) } +func (s *OpenAIGatewayService) handleCodexModelsManifestAccountAuthError(ctx context.Context, account *Account, request codexModelsManifestRequest, err error) { + if s == nil || account == nil || err == nil || request.useAPIKeyUpstream || !account.IsOpenAIOAuth() || account.IsOpenAIAgentIdentity() { + return + } + var upstreamErr *codexModelsManifestUpstreamError + if !errors.As(err, &upstreamErr) || upstreamErr.statusCode != http.StatusUnauthorized { + return + } + headers := upstreamErr.headers + if headers == nil { + headers = http.Header{} + } + s.handleOpenAIAccountUpstreamErrorForModel(ctx, account, "", upstreamErr.statusCode, headers, upstreamErr.body) +} + func (s *OpenAIGatewayService) fetchCachedAPIKeyCodexModelsManifest(ctx context.Context, request codexModelsManifestRequest, ifNoneMatch string) (*CodexModelsManifest, error) { if err := ctx.Err(); err != nil { return nil, err @@ -435,8 +454,10 @@ func (s *OpenAIGatewayService) fetchCodexModelsManifestUpstream(ctx context.Cont return nil, &codexModelsManifestUpstreamError{ err: infraerrors.Newf(http.StatusBadGateway, "OPENAI_CODEX_MODELS_UPSTREAM_FAILED", "codex models manifest upstream error %d: %s", resp.StatusCode, message), statusCode: resp.StatusCode, + headers: resp.Header.Clone(), body: body, - retryable: resp.StatusCode == http.StatusTooManyRequests || + retryable: (resp.StatusCode == http.StatusUnauthorized && !request.useAPIKeyUpstream) || + resp.StatusCode == http.StatusTooManyRequests || (resp.StatusCode >= http.StatusInternalServerError && resp.StatusCode < 600), } } @@ -448,9 +469,84 @@ func (s *OpenAIGatewayService) fetchCodexModelsManifestUpstream(ctx context.Cont retryable: !errors.Is(err, ErrUpstreamResponseBodyTooLarge) && isRetryableCodexModelsManifestTransportError(err), } } + if request.useAPIKeyUpstream { + body = convertOpenAIModelListToCodexManifest(body) + } + if err := validateCodexModelsManifestEnvelope(body); err != nil { + return nil, &codexModelsManifestUpstreamError{ + err: infraerrors.Newf( + http.StatusBadGateway, + "OPENAI_CODEX_MODELS_UPSTREAM_INVALID_MANIFEST", + "codex models manifest upstream returned an invalid envelope: %v", + err, + ), + retryable: true, + } + } return &CodexModelsManifest{Body: body, ETag: resp.Header.Get("ETag")}, nil } +func convertOpenAIModelListToCodexManifest(body []byte) []byte { + var envelope map[string]json.RawMessage + if err := json.Unmarshal(body, &envelope); err != nil || envelope == nil { + return body + } + if _, ok := envelope["models"]; ok { + return body + } + data, ok := envelope["data"] + if !ok { + return body + } + var entries []struct { + ID string `json:"id"` + } + if err := json.Unmarshal(data, &entries); err != nil { + return body + } + type codexModelEntry struct { + Slug string `json:"slug"` + } + models := make([]codexModelEntry, 0, len(entries)) + for _, entry := range entries { + id := strings.TrimSpace(entry.ID) + if id != "" { + models = append(models, codexModelEntry{Slug: id}) + } + } + if len(models) == 0 { + return body + } + converted, err := json.Marshal(map[string][]codexModelEntry{"models": models}) + if err != nil { + return body + } + return converted +} + +func validateCodexModelsManifestEnvelope(body []byte) error { + var envelope map[string]json.RawMessage + if err := json.Unmarshal(body, &envelope); err != nil { + return fmt.Errorf("decode JSON object: %w", err) + } + if envelope == nil { + return errors.New("expected a JSON object") + } + models, ok := envelope["models"] + if !ok { + return errors.New("missing top-level models array") + } + models = bytes.TrimSpace(models) + if len(models) == 0 || models[0] != '[' { + return errors.New("top-level models field is not an array") + } + var entries []json.RawMessage + if err := json.Unmarshal(models, &entries); err != nil { + return fmt.Errorf("decode top-level models array: %w", err) + } + return nil +} + func buildCodexModelsManifestCacheKey(request codexModelsManifestRequest) string { hasher := sha256.New() _, _ = fmt.Fprintf(hasher, "%d\n%s\n%s\n", request.accountID, request.proxyURL, request.url) diff --git a/backend/internal/service/openai_codex_models_service_test.go b/backend/internal/service/openai_codex_models_service_test.go index 33793fdec..c8f60a508 100644 --- a/backend/internal/service/openai_codex_models_service_test.go +++ b/backend/internal/service/openai_codex_models_service_test.go @@ -278,7 +278,7 @@ func TestFetchCodexModelsManifestAPIKeyServesStaleWhileRefreshing(t *testing.T) header := make(http.Header) if call == 1 { header.Set("ETag", `"first"`) - return &http.Response{StatusCode: http.StatusOK, Header: header, Body: io.NopCloser(strings.NewReader(`{"version":1}`))}, nil + return &http.Response{StatusCode: http.StatusOK, Header: header, Body: io.NopCloser(strings.NewReader(`{"models":[],"version":1}`))}, nil } require.Equal(t, `"first"`, req.Header.Get("If-None-Match")) if call == 2 { @@ -286,14 +286,14 @@ func TestFetchCodexModelsManifestAPIKeyServesStaleWhileRefreshing(t *testing.T) <-releaseRefresh } header.Set("ETag", `"second"`) - return &http.Response{StatusCode: http.StatusOK, Header: header, Body: io.NopCloser(strings.NewReader(`{"version":2}`))}, nil + return &http.Response{StatusCode: http.StatusOK, Header: header, Body: io.NopCloser(strings.NewReader(`{"models":[],"version":2}`))}, nil }} s := newCodexModelsAPIKeyTestService(upstream) account := newCodexModelsAPIKeyTestAccount("https://upstream.example/v1") first, err := s.FetchCodexModelsManifest(context.Background(), account, "0.144.0", "") require.NoError(t, err) - require.JSONEq(t, `{"version":1}`, string(first.Body)) + require.JSONEq(t, `{"models":[],"version":1}`, string(first.Body)) s.codexModelsManifestCache.mu.Lock() for key, entry := range s.codexModelsManifestCache.entries { @@ -305,13 +305,13 @@ func TestFetchCodexModelsManifestAPIKeyServesStaleWhileRefreshing(t *testing.T) stale, err := s.FetchCodexModelsManifest(context.Background(), account, "0.144.0", "") require.NoError(t, err) - require.JSONEq(t, `{"version":1}`, string(stale.Body)) + require.JSONEq(t, `{"models":[],"version":1}`, string(stale.Body)) <-refreshStarted close(releaseRefresh) require.Eventually(t, func() bool { refreshed, fetchErr := s.FetchCodexModelsManifest(context.Background(), account, "0.144.0", "") - return fetchErr == nil && string(refreshed.Body) == `{"version":2}` + return fetchErr == nil && string(refreshed.Body) == `{"models":[],"version":2}` }, 2*time.Second, 10*time.Millisecond) require.Equal(t, int32(2), calls.Load()) } @@ -320,7 +320,7 @@ func TestFetchCodexModelsManifestCacheIsolationAndBodyLimits(t *testing.T) { var calls atomic.Int32 upstream := &codexModelsHTTPUpstreamStub{do: func(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) { call := calls.Add(1) - body := `{"call":` + strconv.Itoa(int(call)) + `}` + body := `{"models":[],"call":` + strconv.Itoa(int(call)) + `}` return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(body))}, nil }} s := newCodexModelsAPIKeyTestService(upstream) @@ -431,7 +431,7 @@ func TestFetchCodexModelsManifestCacheKeyIsolatesRequestIdentity(t *testing.T) { func TestFetchCodexModelsManifestDoesNotCacheBodiesOverOneMiB(t *testing.T) { var calls atomic.Int32 - body := strings.Repeat("x", codexModelsManifestCacheBodyLimit+1) + body := `{"models":[],"padding":"` + strings.Repeat("x", codexModelsManifestCacheBodyLimit+1) + `"}` upstream := &codexModelsHTTPUpstreamStub{do: func(_ *http.Request, _ string, _ int64, _ int) (*http.Response, error) { calls.Add(1) return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(body))}, nil @@ -446,6 +446,96 @@ func TestFetchCodexModelsManifestDoesNotCacheBodiesOverOneMiB(t *testing.T) { require.Equal(t, int32(2), calls.Load()) } +func TestFetchCodexModelsManifestRejectsInvalidEnvelopeWithoutCaching(t *testing.T) { + var calls atomic.Int32 + upstream := &codexModelsHTTPUpstreamStub{do: func(_ *http.Request, _ string, _ int64, _ int) (*http.Response, error) { + calls.Add(1) + return &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader(`{"object":"unexpected"}`)), + }, nil + }} + s := newCodexModelsAPIKeyTestService(upstream) + account := newCodexModelsAPIKeyTestAccount("https://upstream.example/v1") + + for range 2 { + _, err := s.FetchCodexModelsManifest(context.Background(), account, "0.144.0", "") + require.Error(t, err) + require.True(t, IsRetryableCodexModelsManifestError(err)) + require.ErrorContains(t, err, "OPENAI_CODEX_MODELS_UPSTREAM_INVALID_MANIFEST") + } + require.Equal(t, int32(2), calls.Load(), "invalid manifests must not enter the cache") +} + +func TestConvertOpenAIModelListToCodexManifest(t *testing.T) { + converted := convertOpenAIModelListToCodexManifest([]byte(`{"object":"list","data":[{"id":"gpt-5.6"},{"id":" "},{"id":"gpt-image-2"}]}`)) + require.JSONEq(t, `{"models":[{"slug":"gpt-5.6"},{"slug":"gpt-image-2"}]}`, string(converted)) + require.NoError(t, validateCodexModelsManifestEnvelope(converted)) +} + +func TestFetchCodexModelsManifestOAuth401IsRetryable(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte(`{"error":{"code":"token_revoked","message":"revoked"}}`)) + })) + defer server.Close() + original := chatgptCodexModelsURL + chatgptCodexModelsURL = server.URL + t.Cleanup(func() { chatgptCodexModelsURL = original }) + + _, err := (&OpenAIGatewayService{}).FetchCodexModelsManifest(context.Background(), newCodexModelsOAuthTestAccount(), "0.144.0", "") + require.Error(t, err) + require.True(t, IsRetryableCodexModelsManifestError(err)) +} + +type codexModelsCanceledStateRepo struct { + AccountRepository + updateCredentialsCalls int + tempUnschedulableCalls int + updateCredentialsCtxErr error + tempUnschedulableCtxErr error +} + +func (r *codexModelsCanceledStateRepo) UpdateCredentials(ctx context.Context, _ int64, _ map[string]any) error { + r.updateCredentialsCalls++ + r.updateCredentialsCtxErr = ctx.Err() + return nil +} + +func (r *codexModelsCanceledStateRepo) SetTempUnschedulable(ctx context.Context, _ int64, _ time.Time, _ string) error { + r.tempUnschedulableCalls++ + r.tempUnschedulableCtxErr = ctx.Err() + return nil +} + +func TestCodexModelsOAuth401PersistsStateAfterCallerCancellation(t *testing.T) { + repo := &codexModelsCanceledStateRepo{} + rateLimitService := NewRateLimitService(repo, nil, &config.Config{}, nil, nil) + svc := &OpenAIGatewayService{rateLimitService: rateLimitService} + account := newCodexModelsOAuthTestAccount() + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + svc.handleCodexModelsManifestAccountAuthError( + ctx, + account, + codexModelsManifestRequest{}, + &codexModelsManifestUpstreamError{ + err: errors.New("models unauthorized"), + statusCode: http.StatusUnauthorized, + headers: http.Header{}, + body: []byte(`{"error":{"message":"expired access token"}}`), + }, + ) + + require.Equal(t, 1, repo.updateCredentialsCalls) + require.NoError(t, repo.updateCredentialsCtxErr) + require.Equal(t, 1, repo.tempUnschedulableCalls) + require.NoError(t, repo.tempUnschedulableCtxErr) + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) +} + func TestFetchCodexModelsManifestSharedRefreshSurvivesCallerCancellation(t *testing.T) { var calls atomic.Int32 readStarted := make(chan struct{}) diff --git a/backend/internal/service/openai_codex_transform.go b/backend/internal/service/openai_codex_transform.go index f253b1c2d..ec6cd39a0 100644 --- a/backend/internal/service/openai_codex_transform.go +++ b/backend/internal/service/openai_codex_transform.go @@ -1,11 +1,37 @@ package service import ( + "crypto/sha256" + "encoding/hex" "encoding/json" "fmt" "strings" ) +const ( + codexCallIDMaxLength = 64 + codexCallIDPrefix = "fc_" +) + +func normalizeCodexCallID(id string) string { + candidate := id + switch { + case id == "": + return "" + case strings.HasPrefix(id, "fc"): + case strings.HasPrefix(id, "call_"): + candidate = codexCallIDPrefix + strings.TrimPrefix(id, "call_") + default: + candidate = codexCallIDPrefix + id + } + if len(candidate) <= codexCallIDMaxLength { + return candidate + } + digest := sha256.Sum256([]byte("sub2api:codex-call-id:v1:" + candidate)) + encoded := hex.EncodeToString(digest[:]) + return codexCallIDPrefix + encoded[:codexCallIDMaxLength-len(codexCallIDPrefix)] +} + var codexModelMap = map[string]string{ "gpt-5.6-sol": "gpt-5.6-sol", "gpt-5.6-terra": "gpt-5.6-terra", @@ -1145,13 +1171,7 @@ func filterCodexInput(input []any, preserveReferences bool) []any { // 仅修正真正的 tool/function call 标识，避免误改普通 message/reasoning id； // 若 item_reference 指向 legacy call_* 标识，则仅修正该引用本身。 fixCallIDPrefix := func(id string) string { - if id == "" || strings.HasPrefix(id, "fc") { - return id - } - if strings.HasPrefix(id, "call_") { - return "fc" + strings.TrimPrefix(id, "call_") - } - return "fc_" + id + return normalizeCodexCallID(id) } if typ == "item_reference" { diff --git a/backend/internal/service/openai_codex_transform_test.go b/backend/internal/service/openai_codex_transform_test.go index be1d37859..2f706f205 100644 --- a/backend/internal/service/openai_codex_transform_test.go +++ b/backend/internal/service/openai_codex_transform_test.go @@ -41,7 +41,7 @@ func TestApplyCodexOAuthTransform_ToolContinuationPreservesInput(t *testing.T) { second, ok := input[1].(map[string]any) require.True(t, ok) require.Equal(t, "o1", second["id"]) - require.Equal(t, "fc1", second["call_id"]) + require.Equal(t, "fc_1", second["call_id"]) } func TestApplyCodexOAuthTransform_ToolContinuationPreservesNativeMessageAndReasoningIDs(t *testing.T) { @@ -87,11 +87,31 @@ func TestApplyCodexOAuthTransform_ToolContinuationNormalizesToolReferenceIDsOnly first, ok := input[0].(map[string]any) require.True(t, ok) - require.Equal(t, "fc1", first["id"]) + require.Equal(t, "fc_1", first["id"]) second, ok := input[1].(map[string]any) require.True(t, ok) - require.Equal(t, "fc1", second["call_id"]) + require.Equal(t, "fc_1", second["call_id"]) +} + +func TestApplyCodexOAuthTransform_BoundsLongCallIDsAndPreservesPairing(t *testing.T) { + suffix := strings.Repeat("z", 62) + reqBody := map[string]any{ + "model": "gpt-5.2", + "input": []any{ + map[string]any{"type": "function_call", "call_id": "call_" + suffix, "name": "shell"}, + map[string]any{"type": "function_call_output", "call_id": "fc_" + suffix, "output": "done"}, + }, + } + + applyCodexOAuthTransform(reqBody, false, false) + + input := reqBody["input"].([]any) + callID := input[0].(map[string]any)["call_id"].(string) + outputCallID := input[1].(map[string]any)["call_id"].(string) + require.LessOrEqual(t, len(callID), codexCallIDMaxLength) + require.True(t, strings.HasPrefix(callID, codexCallIDPrefix)) + require.Equal(t, callID, outputCallID) } func TestApplyCodexOAuthTransform_ToolSearchOutputPreservesCallID(t *testing.T) { @@ -111,7 +131,7 @@ func TestApplyCodexOAuthTransform_ToolSearchOutputPreservesCallID(t *testing.T) first, ok := input[0].(map[string]any) require.True(t, ok) require.Equal(t, "tool_search_output", first["type"]) - require.Equal(t, "fc1", first["call_id"]) + require.Equal(t, "fc_1", first["call_id"]) } func TestApplyCodexOAuthTransform_CustomAndMCPToolOutputsPreserveCallID(t *testing.T) { @@ -131,11 +151,11 @@ func TestApplyCodexOAuthTransform_CustomAndMCPToolOutputsPreserveCallID(t *testi first, ok := input[0].(map[string]any) require.True(t, ok) - require.Equal(t, "fccustom", first["call_id"]) + require.Equal(t, "fc_custom", first["call_id"]) second, ok := input[1].(map[string]any) require.True(t, ok) - require.Equal(t, "fcmcp", second["call_id"]) + require.Equal(t, "fc_mcp", second["call_id"]) } func TestApplyCodexOAuthTransform_ImageAndWebSearchCallsDoNotGainCallID(t *testing.T) { @@ -188,7 +208,7 @@ func TestApplyCodexOAuthTransform_ConvertsToolRoleMessageToFunctionCallOutput(t item, ok := input[0].(map[string]any) require.True(t, ok) require.Equal(t, "function_call_output", item["type"]) - require.Equal(t, "fc1", item["call_id"]) + require.Equal(t, "fc_1", item["call_id"]) require.Equal(t, "ok", item["output"]) _, hasRole := item["role"] require.False(t, hasRole) @@ -307,7 +327,7 @@ func TestApplyCodexOAuthTransform_AddsFallbackNameForFunctionCallInput(t *testin require.True(t, ok) require.Equal(t, "function_call", item["type"]) require.Equal(t, "tool", item["name"]) - require.Equal(t, "fc1", item["call_id"]) + require.Equal(t, "fc_1", item["call_id"]) } func TestApplyCodexOAuthTransform_PreservesFunctionCallInputName(t *testing.T) { @@ -326,7 +346,7 @@ func TestApplyCodexOAuthTransform_PreservesFunctionCallInputName(t *testing.T) { item, ok := input[0].(map[string]any) require.True(t, ok) require.Equal(t, "shell", item["name"]) - require.Equal(t, "fc1", item["call_id"]) + require.Equal(t, "fc_1", item["call_id"]) } func TestApplyCodexOAuthTransform_PreservesMCPToolCallIDAndName(t *testing.T) { @@ -351,7 +371,7 @@ func TestApplyCodexOAuthTransform_PreservesMCPToolCallIDAndName(t *testing.T) { require.True(t, ok) require.Equal(t, "mcp_tool_call", item["type"]) require.Equal(t, "remote_tool", item["name"]) - require.Equal(t, "fcabc", item["call_id"]) + require.Equal(t, "fc_abc", item["call_id"]) } func TestCodexInputItemRequiresNameTypesAllowCallID(t *testing.T) { diff --git a/backend/internal/service/openai_fast_policy_ws_test.go b/backend/internal/service/openai_fast_policy_ws_test.go index 3316a242c..798a5fc12 100644 --- a/backend/internal/service/openai_fast_policy_ws_test.go +++ b/backend/internal/service/openai_fast_policy_ws_test.go @@ -7,6 +7,7 @@ import ( "net/http" "net/http/httptest" "strings" + "sync" "sync/atomic" "testing" "time" @@ -310,6 +311,82 @@ func TestPolicyEnforcingFrameConn_WithoutCapturedFallbackPolicyMisses(t *testing "sanity: without capturedSessionModel fallback the leak (D5) reproduces — confirms the fix is load-bearing") } +func TestPolicyEnforcingFrameConn_RejectsMalformedJSONBeforeFilter(t *testing.T) { + t.Parallel() + + inner := &fakePassthroughFrameConn{reads: [][]byte{[]byte(`{"type":"response.create"`)}} + filterCalled := false + wrapper := &openAIWSPolicyEnforcingFrameConn{ + inner: inner, + filter: func(_ coderws.MessageType, payload []byte) ([]byte, *OpenAIFastBlockedError, error) { + filterCalled = true + return payload, nil, nil + }, + } + + _, payload, err := wrapper.ReadFrame(context.Background()) + require.Error(t, err) + require.Nil(t, payload) + require.False(t, filterCalled, "malformed JSON must not enter policy or reach upstream") + var closeErr *OpenAIWSClientCloseError + require.ErrorAs(t, err, &closeErr) + require.Equal(t, coderws.StatusPolicyViolation, closeErr.StatusCode()) +} + +func TestOpenAIWSPassthroughTurnLifecycleSerializesAndPairsHooks(t *testing.T) { + t.Parallel() + + var ( + mu sync.Mutex + beforeTurns []int + afterTurns []int + afterErrors []error + ) + lifecycle := newOpenAIWSPassthroughTurnLifecycle(&OpenAIWSIngressHooks{ + BeforeTurn: func(turn int) error { + mu.Lock() + beforeTurns = append(beforeTurns, turn) + mu.Unlock() + return nil + }, + AfterTurn: func(turn int, _ *OpenAIForwardResult, turnErr error) { + mu.Lock() + afterTurns = append(afterTurns, turn) + afterErrors = append(afterErrors, turnErr) + mu.Unlock() + }, + }) + + turn, err := lifecycle.begin() + require.NoError(t, err) + require.Equal(t, 1, turn) + _, err = lifecycle.begin() + require.Error(t, err, "a second response.create before terminal must be rejected") + var closeErr *OpenAIWSClientCloseError + require.ErrorAs(t, err, &closeErr) + require.Equal(t, coderws.StatusPolicyViolation, closeErr.StatusCode()) + + finishedTurn, finished := lifecycle.finish(&OpenAIForwardResult{RequestID: "resp_1"}, nil) + require.True(t, finished) + require.Equal(t, 1, finishedTurn) + turn, err = lifecycle.begin() + require.NoError(t, err) + require.Equal(t, 2, turn) + turnErr := errors.New("relay failed") + finishedTurn, finished = lifecycle.finish(nil, turnErr) + require.True(t, finished) + require.Equal(t, 2, finishedTurn) + _, finished = lifecycle.finish(nil, errors.New("duplicate")) + require.False(t, finished, "each active turn must release exactly once") + + mu.Lock() + defer mu.Unlock() + require.Equal(t, []int{1, 2}, beforeTurns) + require.Equal(t, []int{1, 2}, afterTurns) + require.NoError(t, afterErrors[0]) + require.ErrorIs(t, afterErrors[1], turnErr) +} + // --- Ingress end-to-end test (filter path) --- // TestWSResponseCreate_IngressFiltersServiceTierBeforeUpstream wires up the diff --git a/backend/internal/service/openai_first_output_timeout.go b/backend/internal/service/openai_first_output_timeout.go index dda5f5977..e2a82f6ba 100644 --- a/backend/internal/service/openai_first_output_timeout.go +++ b/backend/internal/service/openai_first_output_timeout.go @@ -25,13 +25,156 @@ const ( openAIFirstOutputScannerFramingAllowance = 64 openAIFirstOutputGuardQueueSize = 1 openAIDefaultStreamQueueSize = 16 + openAIFirstOutputPhaseRoutingBudget = "routing_budget" + openAIFirstOutputPhaseResponseHeaders = "response_headers" + openAIFirstOutputPhaseSemanticOutput = "semantic_output" ) var ( - errOpenAIFirstOutputStageLimit = errors.New("openai first-output staging limit exceeded") - errOpenAIFirstOutputScannerLimit = errors.New("openai pre-output scanner token limit exceeded") + errOpenAIFirstOutputStageLimit = errors.New("openai first-output staging limit exceeded") + errOpenAIFirstOutputScannerLimit = errors.New("openai pre-output scanner token limit exceeded") + ErrOpenAIFirstOutputRoutingBudgetExceeded = errors.New("openai first-output routing budget exceeded") ) +// openAIFirstOutputStartContextKey carries the end-to-end gateway routing start +// time from the HTTP handler into every upstream retry. ForwardWithAnalysis is +// called after account selection, so using time.Now() there would exclude +// queueing and account selection from the first-output budget. +type openAIFirstOutputStartContextKey struct{} + +type openAIFirstOutputBudgetContextKey struct{} + +type openAIFirstOutputBudgetContextValue struct { + enabled bool + deadline time.Time +} + +// WithOpenAIFirstOutputStart attaches the routing start time used by the +// native OpenAI Responses first-output guard. A zero value is ignored so a +// caller cannot accidentally replace a valid start time with an invalid one. +func WithOpenAIFirstOutputStart(ctx context.Context, start time.Time) context.Context { + if ctx == nil || start.IsZero() { + return ctx + } + return context.WithValue(ctx, openAIFirstOutputStartContextKey{}, start) +} + +func openAIFirstOutputStart(ctx context.Context) time.Time { + if ctx != nil { + if start, ok := ctx.Value(openAIFirstOutputStartContextKey{}).(time.Time); ok && !start.IsZero() { + return start + } + } + return time.Now() +} + +// ensureOpenAIFirstOutputStart makes the fallback start time stable for +// service callers that invoke ForwardWithAnalysis directly instead of going +// through an HTTP handler. Without this, the header guard and stream timer +// could each choose a different "now" and silently extend the budget. +func ensureOpenAIFirstOutputStart(ctx context.Context) (context.Context, time.Time) { + if ctx == nil { + ctx = context.Background() + } + if ctx != nil { + if start, ok := ctx.Value(openAIFirstOutputStartContextKey{}).(time.Time); ok && !start.IsZero() { + return ctx, start + } + } + start := time.Now() + return WithOpenAIFirstOutputStart(ctx, start), start +} + +// WithOpenAIFirstOutputBudget stores an absolute, request-scoped routing +// budget without putting a deadline on the request context itself. The latter +// is intentional: once semantic output starts, the normal long-lived stream +// must not be cancelled by the pre-output budget. A non-positive timeout +// explicitly disables the budget and shadows any inherited upper bound. +func WithOpenAIFirstOutputBudget(ctx context.Context, timeout time.Duration) context.Context { + if ctx == nil { + ctx = context.Background() + } + value := openAIFirstOutputBudgetContextValue{} + if timeout > 0 { + value.enabled = true + value.deadline = openAIFirstOutputStart(ctx).Add(timeout) + } + return context.WithValue(ctx, openAIFirstOutputBudgetContextKey{}, value) +} + +// OpenAIFirstOutputBudgetRemaining returns the remaining routing budget and +// whether a budget is enabled. A negative/zero duration means the budget has +// expired; callers should fail fast instead of starting another wait/dial. +func OpenAIFirstOutputBudgetRemaining(ctx context.Context) (time.Duration, bool) { + if ctx == nil { + return 0, false + } + value, ok := ctx.Value(openAIFirstOutputBudgetContextKey{}).(openAIFirstOutputBudgetContextValue) + if !ok || !value.enabled || value.deadline.IsZero() { + return 0, false + } + return time.Until(value.deadline), true +} + +// CapOpenAIFirstOutputWait applies the remaining routing budget to a local +// wait timeout. It never changes the caller's context or affects streaming +// after the first semantic output. +func CapOpenAIFirstOutputWait(ctx context.Context, requested time.Duration) time.Duration { + remaining, enabled := OpenAIFirstOutputBudgetRemaining(ctx) + if !enabled || requested <= 0 { + return requested + } + if remaining <= 0 { + return time.Nanosecond + } + if remaining < requested { + return remaining + } + return requested +} + +// WithOpenAIFirstOutputRoutingDeadline converts the soft request budget into a +// cancellable child context for routing-only work (moderation, cache/DB lookup, +// and account selection). Callers must cancel it before starting the upstream +// request so the deadline can never terminate an already-producing stream. +func WithOpenAIFirstOutputRoutingDeadline(ctx context.Context) (context.Context, context.CancelFunc) { + if ctx == nil { + ctx = context.Background() + } + remaining, enabled := OpenAIFirstOutputBudgetRemaining(ctx) + if !enabled { + return ctx, func() {} + } + if remaining <= 0 { + remaining = time.Nanosecond + } + return context.WithTimeoutCause(ctx, remaining, ErrOpenAIFirstOutputRoutingBudgetExceeded) +} + +func (s *OpenAIGatewayService) OpenAIFirstOutputRoutingBudget(body []byte, modelCandidates ...string) time.Duration { + if s == nil || s.cfg == nil { + return 0 + } + effort := extractOpenAIReasoningEffortFromBody(body, modelCandidates...) + effortValue := "" + if effort != nil { + effortValue = *effort + } + return s.openAIFirstOutputTimeout(effortValue) +} + +func (s *OpenAIGatewayService) OpenAIFirstOutputBudgetForAccount(account *Account, body []byte, modelCandidates ...string) time.Duration { + if s == nil || account == nil || account.Platform != PlatformOpenAI { + return 0 + } + effort := extractOpenAIReasoningEffortFromBody(body, modelCandidates...) + effortValue := "" + if effort != nil { + effortValue = *effort + } + return s.openAIFirstOutputTimeout(effortValue) +} + type openAIFirstOutputStage struct { limit int64 size int64 @@ -255,20 +398,44 @@ func (s *OpenAIGatewayService) newOpenAIFirstOutputTimeoutError( account.ID, originalModel, reasoningEffort, phase, elapsed, timeout, ) requestID := strings.TrimSpace(responseHeaders.Get("x-request-id")) + eventMessage := "OpenAI upstream produced no semantic output before the deadline" + responseType := "first_output_timeout" + responseMessage := "Upstream produced no output before the deadline" + if phase == openAIFirstOutputPhaseRoutingBudget { + eventMessage = "OpenAI gateway routing budget expired before an upstream attempt" + responseType = "routing_budget_exhausted" + responseMessage = "Gateway routing budget expired before an upstream attempt could start" + } appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ Platform: account.Platform, AccountID: account.ID, AccountName: account.Name, UpstreamStatusCode: http.StatusGatewayTimeout, UpstreamRequestID: requestID, - Kind: "first_output_timeout", Message: "OpenAI upstream produced no semantic output before the deadline", + Kind: responseType, Message: eventMessage, Detail: fmt.Sprintf("phase=%s elapsed_ms=%d timeout_ms=%d", phase, elapsed.Milliseconds(), timeout.Milliseconds()), }) - if s.rateLimitService != nil { - s.rateLimitService.HandleStreamTimeout(ctx, account, originalModel) + // Do not call RateLimitService.HandleStreamTimeout here. That method models + // an idle stream-data timeout and may synchronously update persistent + // account state; a first-output deadline is a separate signal. The handler + // still reports genuine upstream failures to the scheduler's EWMA/error + // stats, while routing-budget exhaustion is explicitly excluded below. + reason := GatewayFailureReasonOpenAIFirstOutputTimeout + scope := GatewayFailureScopeProvider + nextAccountAction := NextAccountLegacyRetry + if phase == openAIFirstOutputPhaseRoutingBudget { + reason = GatewayFailureReasonRoutingBudgetExhausted + scope = GatewayFailureScopeRequest + // No account can recover an already exhausted absolute budget. Stop here + // so account-switch metrics and selection work are not polluted by a + // retry that cannot reach the upstream. + nextAccountAction = NextAccountStop } return &UpstreamFailoverError{ StatusCode: http.StatusGatewayTimeout, - ResponseBody: []byte(`{"error":{"type":"first_output_timeout","message":"Upstream produced no output before the deadline"}}`), + ResponseBody: []byte(fmt.Sprintf(`{"error":{"type":%q,"message":%q}}`, responseType, responseMessage)), ResponseHeaders: responseHeaders.Clone(), SafeToFailoverAfterWrite: true, + Scope: scope, + Reason: reason, + NextAccountAction: nextAccountAction, } } diff --git a/backend/internal/service/openai_first_output_timeout_test.go b/backend/internal/service/openai_first_output_timeout_test.go index a9acf2a9e..afe75f9f2 100644 --- a/backend/internal/service/openai_first_output_timeout_test.go +++ b/backend/internal/service/openai_first_output_timeout_test.go @@ -2,13 +2,54 @@ package service import ( "bytes" + "context" "errors" "os" "testing" + "time" "github.com/stretchr/testify/require" ) +func TestOpenAIFirstOutputStartPreservesEndToEndRequestStart(t *testing.T) { + start := time.Now().Add(-2 * time.Second) + ctx := WithOpenAIFirstOutputStart(context.Background(), start) + recovered := openAIFirstOutputStart(ctx) + require.Equal(t, start, recovered) + + // Contexts without the marker retain the compatibility behavior and start + // timing at the service entry point. + compatStart := openAIFirstOutputStart(context.Background()) + require.WithinDuration(t, time.Now(), compatStart, time.Second) +} + +func TestEnsureOpenAIFirstOutputStartIsStableForDirectServiceCallers(t *testing.T) { + ctx, first := ensureOpenAIFirstOutputStart(context.Background()) + recovered := openAIFirstOutputStart(ctx) + require.Equal(t, first, recovered) + + ctx2, second := ensureOpenAIFirstOutputStart(ctx) + require.Same(t, ctx, ctx2) + require.Equal(t, first, second) +} + +func TestOpenAIFirstOutputRoutingBudgetDoesNotPenalizeAccount(t *testing.T) { + err := (&OpenAIGatewayService{}).newOpenAIFirstOutputTimeoutError( + context.Background(), + nil, + &Account{ID: 1, Platform: PlatformOpenAI}, + time.Now().Add(-time.Second), + "gpt-5", + "", + time.Second, + openAIFirstOutputPhaseRoutingBudget, + nil, + ) + require.Equal(t, GatewayFailureReasonRoutingBudgetExhausted, err.Reason) + require.False(t, err.ShouldReportAccountScheduleFailure()) + require.False(t, err.ShouldRetryNextAccount()) +} + func TestOpenAIFirstOutputStageUnlinkFailureFailsFastAndRetriesCleanup(t *testing.T) { stage := newDefaultOpenAIFirstOutputStage() stage.memoryOnly = false diff --git a/backend/internal/service/openai_gateway_chat_completions.go b/backend/internal/service/openai_gateway_chat_completions.go index 3bc8e092c..58c50289f 100644 --- a/backend/internal/service/openai_gateway_chat_completions.go +++ b/backend/internal/service/openai_gateway_chat_completions.go @@ -270,14 +270,14 @@ func (s *OpenAIGatewayService) ForwardAsChatCompletions( Message: upstreamMsg, Detail: upstreamDetail, }) - if s.rateLimitService != nil { - s.rateLimitService.HandleUpstreamErrorForModel(ctx, account, originalModel, resp.StatusCode, resp.Header, respBody) - } - return nil, &UpstreamFailoverError{ - StatusCode: resp.StatusCode, - ResponseBody: respBody, - RetryableOnSameAccount: shouldRetryOpenAIOnSamePoolAccount(account, resp.StatusCode, upstreamMsg, respBody), - } + s.handleOpenAIAccountUpstreamErrorForModel(ctx, account, billingModel, resp.StatusCode, resp.Header, respBody) + return nil, newOpenAIUpstreamFailoverError( + resp.StatusCode, + resp.Header, + respBody, + upstreamMsg, + shouldRetryOpenAIOnSamePoolAccount(account, resp.StatusCode, upstreamMsg, respBody), + ) } return s.handleChatCompletionsErrorResponse(resp, c, account, originalModel) } diff --git a/backend/internal/service/openai_gateway_chat_completions_raw.go b/backend/internal/service/openai_gateway_chat_completions_raw.go index d6071d983..09b142b21 100644 --- a/backend/internal/service/openai_gateway_chat_completions_raw.go +++ b/backend/internal/service/openai_gateway_chat_completions_raw.go @@ -146,14 +146,14 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions( Kind: "failover", Message: upstreamMsg, }) - if s.rateLimitService != nil { - s.rateLimitService.HandleUpstreamErrorForModel(ctx, account, originalModel, resp.StatusCode, resp.Header, respBody) - } - return nil, &UpstreamFailoverError{ - StatusCode: resp.StatusCode, - ResponseBody: respBody, - RetryableOnSameAccount: shouldRetryOpenAIOnSamePoolAccount(account, resp.StatusCode, upstreamMsg, respBody), - } + s.handleOpenAIAccountUpstreamErrorForModel(ctx, account, originalModel, resp.StatusCode, resp.Header, respBody) + return nil, newOpenAIUpstreamFailoverError( + resp.StatusCode, + resp.Header, + respBody, + upstreamMsg, + shouldRetryOpenAIOnSamePoolAccount(account, resp.StatusCode, upstreamMsg, respBody), + ) } return s.handleChatCompletionsErrorResponse(resp, c, account, originalModel) } diff --git a/backend/internal/service/openai_gateway_chat_completions_test.go b/backend/internal/service/openai_gateway_chat_completions_test.go index bf9b738c7..b09f6de37 100644 --- a/backend/internal/service/openai_gateway_chat_completions_test.go +++ b/backend/internal/service/openai_gateway_chat_completions_test.go @@ -14,11 +14,33 @@ import ( "github.com/Wei-Shaw/sub2api/internal/config" "github.com/Wei-Shaw/sub2api/internal/pkg/apicompat" "github.com/Wei-Shaw/sub2api/internal/pkg/openai_compat" + "github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" "github.com/tidwall/gjson" ) +type forceChatErrorUpstream struct { + statusCode int + headers http.Header + body string + calls int +} + +func (u *forceChatErrorUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) { + u.calls++ + return &http.Response{ + StatusCode: u.statusCode, + Header: u.headers.Clone(), + Body: io.NopCloser(strings.NewReader(u.body)), + Request: req, + }, nil +} + +func (u *forceChatErrorUpstream) DoWithTLS(req *http.Request, proxyURL string, accountID int64, accountConcurrency int, _ *tlsfingerprint.Profile) (*http.Response, error) { + return u.Do(req, proxyURL, accountID, accountConcurrency) +} + func TestForceChatAnthropicDirectBridgeNonStreaming(t *testing.T) { gin.SetMode(gin.TestMode) body := []byte(`{"model":"gpt-5.5","max_tokens":256,"messages":[{"role":"user","content":"hello"}],"tools":[{"name":"Read","input_schema":{"type":"object"}}]}`) @@ -108,6 +130,87 @@ func TestForceChatResponsesFallbackNonStreaming(t *testing.T) { require.Equal(t, 2, result.Usage.CacheCreationInputTokens) } +func TestForceChatResponsesRepeatedTransientErrorsBlockOnlyRequestedModel(t *testing.T) { + gin.SetMode(gin.TestMode) + upstream := &forceChatErrorUpstream{ + statusCode: http.StatusServiceUnavailable, + headers: http.Header{"Content-Type": []string{"application/json"}}, + body: `{"error":{"message":"temporarily unavailable"}}`, + } + svc := newForceChatBridgeTestService(upstream) + account := newForceChatBridgeTestAccount() + body := []byte(`{"model":"gpt-5.5","input":"hello","stream":false}`) + + for range 2 { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body)) + _, err := svc.forwardResponsesViaRawChatCompletions(context.Background(), c, account, body) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + } + + require.Equal(t, 2, upstream.calls) + require.True(t, svc.isOpenAIAccountModelRuntimeBlocked(account, "gpt-5.5")) + require.False(t, svc.isOpenAIAccountModelRuntimeBlocked(account, "gpt-5.6")) +} + +func TestForwardAsChatCompletionsTransientFailureUsesEffectiveFallbackModel(t *testing.T) { + gin.SetMode(gin.TestMode) + upstream := &forceChatErrorUpstream{ + statusCode: http.StatusServiceUnavailable, + headers: http.Header{"Content-Type": []string{"application/json"}}, + body: `{"error":{"message":"temporarily unavailable"}}`, + } + svc := newForceChatBridgeTestService(upstream) + account := newForceChatBridgeTestAccount() + account.Extra[openai_compat.ExtraKeyResponsesSupported] = true + body := []byte(`{"model":"gpt-requested","stream":false,"messages":[{"role":"user","content":"hello"}]}`) + + for range 2 { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body)) + _, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "gpt-fallback") + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + } + + require.True(t, svc.isOpenAIAccountModelRuntimeBlocked(account, "gpt-fallback")) + require.False(t, svc.isOpenAIAccountModelRuntimeBlocked(account, "gpt-requested")) +} + +func TestForceChatResponsesRepeated413RemainAccountScopedAndSanitized(t *testing.T) { + gin.SetMode(gin.TestMode) + upstream := &forceChatErrorUpstream{ + statusCode: http.StatusRequestEntityTooLarge, + headers: http.Header{"X-Request-Id": []string{"req-body-limit"}}, + body: `{"error":{"message":"proxy internal.example rejected tenant-secret payload size"}}`, + } + svc := newForceChatBridgeTestService(upstream) + body := []byte(`{"model":"gpt-5.5","input":"hello","stream":false}`) + + for attempt := range 2 { + account := newForceChatBridgeTestAccount() + account.ID += int64(attempt) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body)) + _, err := svc.forwardResponsesViaRawChatCompletions(context.Background(), c, account, body) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.True(t, failoverErr.IsOpenAIRequestBodyTooLarge()) + require.Equal(t, GatewayFailureScopeAccount, failoverErr.Scope) + require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction) + require.False(t, failoverErr.RetryableOnSameAccount) + require.Equal(t, http.StatusRequestEntityTooLarge, failoverErr.ClientStatusCode) + require.Equal(t, OpenAIRequestBodyTooLargeClientMessage, failoverErr.ClientMessage) + require.NotContains(t, failoverErr.ClientMessage, "internal.example") + require.Equal(t, "req-body-limit", failoverErr.ResponseHeaders.Get("X-Request-Id")) + } + require.Equal(t, 2, upstream.calls) +} + func TestForceChatAnthropicMissingModelUsesAnthropicError(t *testing.T) { gin.SetMode(gin.TestMode) recorder := httptest.NewRecorder() diff --git a/backend/internal/service/openai_gateway_messages.go b/backend/internal/service/openai_gateway_messages.go index ae616a251..d875ef5b2 100644 --- a/backend/internal/service/openai_gateway_messages.go +++ b/backend/internal/service/openai_gateway_messages.go @@ -301,14 +301,14 @@ func (s *OpenAIGatewayService) ForwardAsAnthropic( Message: upstreamMsg, Detail: upstreamDetail, }) - if s.rateLimitService != nil { - s.rateLimitService.HandleUpstreamErrorForModel(ctx, account, originalModel, resp.StatusCode, resp.Header, respBody) - } - return nil, &UpstreamFailoverError{ - StatusCode: resp.StatusCode, - ResponseBody: respBody, - RetryableOnSameAccount: shouldRetryOpenAIOnSamePoolAccount(account, resp.StatusCode, upstreamMsg, respBody), - } + s.handleOpenAIAccountUpstreamErrorForModel(ctx, account, originalModel, resp.StatusCode, resp.Header, respBody) + return nil, newOpenAIUpstreamFailoverError( + resp.StatusCode, + resp.Header, + respBody, + upstreamMsg, + shouldRetryOpenAIOnSamePoolAccount(account, resp.StatusCode, upstreamMsg, respBody), + ) } // Non-failover error: return Anthropic-formatted error to client return s.handleAnthropicErrorResponse(resp, c, account, originalModel) @@ -475,6 +475,7 @@ func (s *OpenAIGatewayService) handleAnthropicBufferedStreamingResponse( if s.responseHeaderFilter != nil { responseheaders.WriteFilteredHeaders(c.Writer.Header(), resp.Header, s.responseHeaderFilter) } + c.Header("Content-Type", "application/json; charset=utf-8") c.JSON(http.StatusOK, anthropicResp) return &OpenAIForwardResult{ diff --git a/backend/internal/service/openai_gateway_record_usage_test.go b/backend/internal/service/openai_gateway_record_usage_test.go index 233801dbb..8caf0f326 100644 --- a/backend/internal/service/openai_gateway_record_usage_test.go +++ b/backend/internal/service/openai_gateway_record_usage_test.go @@ -1767,10 +1767,13 @@ func TestOpenAIGatewayServiceRecordUsage_ImageUsesChannelTokenPricingWithOfficia RequestID: "resp_image_official_token_usage", Model: "gpt-image-2", Usage: OpenAIUsage{ - InputTokens: 22, - TextInputTokens: 22, - OutputTokens: 196, - ImageOutputTokens: 196, + InputTokens: 42, + TextInputTokens: 22, + ImageInputTokens: 20, + CacheReadInputTokens: 10, + ImageCacheReadInputTokens: 10, + OutputTokens: 196, + ImageOutputTokens: 196, }, ImageCount: 1, ImageSize: "1K", @@ -1793,11 +1796,15 @@ func TestOpenAIGatewayServiceRecordUsage_ImageUsesChannelTokenPricingWithOfficia require.NotNil(t, usageRepo.lastLog.BillingMode) require.Equal(t, string(BillingModeToken), *usageRepo.lastLog.BillingMode) require.Equal(t, 1, usageRepo.lastLog.ImageCount) + require.Equal(t, 32, usageRepo.lastLog.InputTokens) + require.Equal(t, 10, usageRepo.lastLog.ImageInputTokens) require.InDelta(t, 22*5e-6, usageRepo.lastLog.InputCost, 1e-12) + require.InDelta(t, 10*8e-6, usageRepo.lastLog.ImageInputCost, 1e-12) + require.InDelta(t, 10*2e-6, usageRepo.lastLog.CacheReadCost, 1e-12) require.InDelta(t, 0.0, usageRepo.lastLog.OutputCost, 1e-12) require.InDelta(t, 196*30e-6, usageRepo.lastLog.ImageOutputCost, 1e-12) - require.InDelta(t, 0.00599, usageRepo.lastLog.TotalCost, 1e-12) - require.InDelta(t, 0.04792, usageRepo.lastLog.ActualCost, 1e-12) + require.InDelta(t, 0.00609, usageRepo.lastLog.TotalCost, 1e-12) + require.InDelta(t, 0.04872, usageRepo.lastLog.ActualCost, 1e-12) } func TestOpenAIGatewayServiceRecordUsage_ResponseImageBillingModelNotOverriddenByRequestedModel(t *testing.T) { diff --git a/backend/internal/service/openai_gateway_responses_chat_fallback.go b/backend/internal/service/openai_gateway_responses_chat_fallback.go index 0d2ecc801..530557ce9 100644 --- a/backend/internal/service/openai_gateway_responses_chat_fallback.go +++ b/backend/internal/service/openai_gateway_responses_chat_fallback.go @@ -201,14 +201,14 @@ func (s *OpenAIGatewayService) forwardResponsesViaRawChatCompletions( Message: upstreamMsg, Detail: upstreamDetail, }) - if s.rateLimitService != nil { - s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) - } - return nil, &UpstreamFailoverError{ - StatusCode: resp.StatusCode, - ResponseBody: respBody, - RetryableOnSameAccount: account.IsPoolMode() && (isPoolModeRetryableStatus(resp.StatusCode) || isOpenAITransientProcessingError(resp.StatusCode, upstreamMsg, respBody)), - } + s.handleOpenAIAccountUpstreamErrorForModel(ctx, account, originalModel, resp.StatusCode, resp.Header, respBody) + return nil, newOpenAIUpstreamFailoverError( + resp.StatusCode, + resp.Header, + respBody, + upstreamMsg, + account.IsPoolMode() && (isPoolModeRetryableStatus(resp.StatusCode) || isOpenAITransientProcessingError(resp.StatusCode, upstreamMsg, respBody)), + ) } return s.handleErrorResponse(ctx, resp, c, account, chatBody, billingModel) } diff --git a/backend/internal/service/openai_gateway_service.go b/backend/internal/service/openai_gateway_service.go index 8cad23fa4..039a2542f 100644 --- a/backend/internal/service/openai_gateway_service.go +++ b/backend/internal/service/openai_gateway_service.go @@ -435,11 +435,13 @@ type OpenAIGatewayService struct { openaiWSStateStoreOnce sync.Once openaiSchedulerOnce sync.Once openaiWSPassthroughDialerOnce sync.Once + openaiModelTransientOnce sync.Once openaiWSPool *openAIWSConnPool openaiWSStateStore OpenAIWSStateStore openaiScheduler OpenAIAccountScheduler openaiWSPassthroughDialer openAIWSClientDialer openaiAccountStats *openAIAccountRuntimeStats + openaiModelTransient *openAIAccountModelTransientState agentIdentityTaskMu sync.Mutex openaiWSFallbackUntil sync.Map // key: int64(accountID), value: time.Time @@ -517,6 +519,7 @@ func NewOpenAIGatewayService( accountShareModeService: accountShareModeService, responseHeaderFilter: compileResponseHeaderFilter(cfg), codexSnapshotThrottle: newAccountWriteThrottle(openAICodexSnapshotPersistMinInterval), + openaiModelTransient: newOpenAIAccountModelTransientState(openAIModelTransientDefaultMax), } svc.logOpenAIWSModeBootstrap() return svc @@ -1499,14 +1502,29 @@ func (s *OpenAIGatewayService) SelectAccountForModelWithExclusions(ctx context.C // noAvailableOpenAISelectionError builds the standard "no account available" error // while preserving the compact-specific error when applicable. +type noAvailableOpenAIAccountSelectionError struct { + message string +} + +func (e *noAvailableOpenAIAccountSelectionError) Error() string { + if e == nil { + return "no available OpenAI accounts" + } + return e.message +} + +func (e *noAvailableOpenAIAccountSelectionError) Unwrap() error { + return ErrNoAvailableAccounts +} + func noAvailableOpenAISelectionError(requestedModel string, compactBlocked bool) error { if compactBlocked { return ErrNoAvailableCompactAccounts } if requestedModel != "" { - return fmt.Errorf("no available OpenAI accounts supporting model: %s", requestedModel) + return &noAvailableOpenAIAccountSelectionError{message: fmt.Sprintf("no available OpenAI accounts supporting model: %s", requestedModel)} } - return errors.New("no available OpenAI accounts") + return &noAvailableOpenAIAccountSelectionError{message: "no available OpenAI accounts"} } // openAICompactSupportTier classifies an OpenAI account by compact capability. @@ -1670,7 +1688,7 @@ func (s *OpenAIGatewayService) tryStickySessionHit(ctx context.Context, groupID } // 楠岃瘉璐﹀彿鏄惁鍙敤浜庡綋鍓嶈姹? // Verify account is usable for current request - if !isOpenAIAccountEligibleForRequest(account, requestedModel, false) { + if !isOpenAIAccountEligibleForRequest(account, requestedModel, false) || s.isOpenAIAccountRequestRuntimeBlocked(account, requestedModel) { return nil } account = s.recheckSelectedOpenAIAccountFromDB(ctx, groupID, account, requestedModel, requireCompact) @@ -1793,6 +1811,14 @@ func (s *OpenAIGatewayService) SelectAccountWithLoadAwareness(ctx context.Contex } func (s *OpenAIGatewayService) selectAccountWithLoadAwareness(ctx context.Context, groupID *int64, sessionHash string, requestedModel string, excludedIDs map[int64]struct{}, requireCompact bool) (*AccountSelectionResult, error) { + // Keep the public legacy selector's historical sticky-wait behavior for + // callers that explicitly use it. HTTP/Chat routing goes through the + // request-aware variant below, which only preserves a sticky wait for a + // continuation that carries previous_response_id (WS v2). + return s.selectAccountWithLoadAwarenessForRequest(ctx, groupID, sessionHash, requestedModel, excludedIDs, requireCompact, true) +} + +func (s *OpenAIGatewayService) selectAccountWithLoadAwarenessForRequest(ctx context.Context, groupID *int64, sessionHash string, requestedModel string, excludedIDs map[int64]struct{}, requireCompact, preserveStickyWait bool) (*AccountSelectionResult, error) { if s.checkChannelPricingRestriction(ctx, groupID, requestedModel) { slog.Warn("channel pricing restriction blocked request", "group_id", derefGroupID(groupID), @@ -1816,7 +1842,7 @@ func (s *OpenAIGatewayService) selectAccountWithLoadAwareness(ctx context.Contex if err == nil && result.Acquired { return newAccountShareModeSelectionResult(account, true, result.ReleaseFunc, nil), nil } - if stickyAccountID > 0 && stickyAccountID == account.ID && s.concurrencyService != nil { + if preserveStickyWait && stickyAccountID > 0 && stickyAccountID == account.ID && s.concurrencyService != nil { waitingCount, _ := s.concurrencyService.GetAccountWaitingCount(ctx, account.ID) if waitingCount < cfg.StickySessionMaxWaiting { return newAccountShareModeSelectionResult(account, false, nil, &AccountWaitPlan{ @@ -1843,7 +1869,7 @@ func (s *OpenAIGatewayService) selectAccountWithLoadAwareness(ctx context.Contex if err == nil && result.Acquired { return s.newSelectionResult(ctx, account, true, result.ReleaseFunc, nil) } - if stickyAccountID > 0 && stickyAccountID == account.ID && s.concurrencyService != nil { + if preserveStickyWait && stickyAccountID > 0 && stickyAccountID == account.ID && s.concurrencyService != nil { waitingCount, _ := s.concurrencyService.GetAccountWaitingCount(ctx, account.ID) if waitingCount < cfg.StickySessionMaxWaiting { return s.newSelectionResult(ctx, account, false, nil, &AccountWaitPlan{ @@ -1888,7 +1914,7 @@ func (s *OpenAIGatewayService) selectAccountWithLoadAwareness(ctx context.Contex if clearSticky { _ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash) } - if !clearSticky && isOpenAIAccountEligibleForRequest(account, requestedModel, false) { + if !clearSticky && isOpenAIAccountEligibleForRequest(account, requestedModel, false) && !s.isOpenAIAccountRequestRuntimeBlocked(account, requestedModel) { account = s.recheckSelectedOpenAIAccountFromDB(ctx, groupID, account, requestedModel, requireCompact) if account == nil { _ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash) @@ -1902,7 +1928,7 @@ func (s *OpenAIGatewayService) selectAccountWithLoadAwareness(ctx context.Contex } waitingCount, _ := s.concurrencyService.GetAccountWaitingCount(ctx, accountID) - if waitingCount < cfg.StickySessionMaxWaiting { + if preserveStickyWait && waitingCount < cfg.StickySessionMaxWaiting { return s.newSelectionResult(ctx, account, false, nil, &AccountWaitPlan{ AccountID: accountID, MaxConcurrency: account.Concurrency, @@ -2147,7 +2173,7 @@ func (s *OpenAIGatewayService) resolveAccountShareModeBoundAccount(ctx context.C if account.IsOpenAICompatible() && account.IsSchedulable() && requestedModel != "" && !account.IsModelSupported(requestedModel) { return nil, true, accountShareModeUnsupportedModelError(requestedModel) } - if !isOpenAIAccountEligibleForRequest(account, requestedModel, requireCompact) { + if !isOpenAIAccountEligibleForRequest(account, requestedModel, requireCompact) || s.isOpenAIAccountRequestRuntimeBlocked(account, requestedModel) { return nil, true, noAvailableOpenAISelectionError(requestedModel, requireCompact && openAICompactSupportTier(account) == 0) } if s.needsUpstreamChannelRestrictionCheck(ctx, groupID) && s.isUpstreamModelRestrictedByChannel(ctx, *groupID, account, requestedModel, requireCompact) { @@ -2177,7 +2203,7 @@ func (s *OpenAIGatewayService) resolveFreshSchedulableOpenAIAccount(ctx context. fresh = current } - if !isOpenAIAccountEligibleForRequest(fresh, requestedModel, requireCompact) { + if !isOpenAIAccountEligibleForRequest(fresh, requestedModel, requireCompact) || s.isOpenAIAccountRequestRuntimeBlocked(fresh, requestedModel) { return nil } if !IsAccountVisibleToRequestUser(ctx, fresh) { @@ -2214,7 +2240,7 @@ func (s *OpenAIGatewayService) recheckSelectedOpenAIAccountFromDB(ctx context.Co return nil } if s.schedulerSnapshot == nil || s.accountRepo == nil { - if !isOpenAIAccountEligibleForRequest(account, requestedModel, requireCompact) { + if !isOpenAIAccountEligibleForRequest(account, requestedModel, requireCompact) || s.isOpenAIAccountRequestRuntimeBlocked(account, requestedModel) { return nil } if !s.isOpenAIAccountInRequestGroup(account, groupID) { @@ -2227,7 +2253,7 @@ func (s *OpenAIGatewayService) recheckSelectedOpenAIAccountFromDB(ctx context.Co if err != nil || latest == nil { return nil } - if !isOpenAIAccountEligibleForRequest(latest, requestedModel, requireCompact) { + if !isOpenAIAccountEligibleForRequest(latest, requestedModel, requireCompact) || s.isOpenAIAccountRequestRuntimeBlocked(latest, requestedModel) { return nil } if !IsAccountVisibleToRequestUser(ctx, latest) { @@ -2239,6 +2265,62 @@ func (s *OpenAIGatewayService) recheckSelectedOpenAIAccountFromDB(ctx context.Co return latest } +// RevalidateSelectedOpenAIAccountForDispatch closes the scheduling-to-dispatch +// race for queued HTTP requests and multi-turn WebSocket sessions. Normal +// groups are authorized from the latest account row and group bindings. Account +// share mode groups instead re-resolve the request's active membership, because +// their listing accounts are intentionally private to the owner. +func (s *OpenAIGatewayService) RevalidateSelectedOpenAIAccountForDispatch( + ctx context.Context, + groupID *int64, + account *Account, + requirements OpenAIAccountDispatchRequirements, +) (*Account, error) { + if s == nil || s.accountRepo == nil || account == nil || account.ID <= 0 { + return nil, ErrNoAvailableAccounts + } + + isModeGroup := groupID != nil && *groupID > 0 && s.accountShareModeService != nil && s.accountShareModeService.IsModeGroup(ctx, *groupID) + if isModeGroup { + requestCtx, ok := AccountShareModeRequestFromContext(ctx) + if !ok { + return nil, ErrAccountShareModeGroupUnbound + } + // Use a fresh request state so a long-lived WebSocket cannot reuse the + // membership cached during its first turn after that membership ends. + freshBindingCtx := WithAccountShareModeRequest(ctx, requestCtx.UserID, requestCtx.APIKeyID) + membership, _, err := s.accountShareModeService.ResolveActiveBindingForRequest(freshBindingCtx, requestCtx.UserID, requestCtx.APIKeyID, *groupID) + if err != nil { + return nil, err + } + if membership == nil || membership.AccountID != account.ID { + return nil, ErrAccountShareModeGroupUnbound + } + } + + latest, err := s.accountRepo.GetByID(ctx, account.ID) + if err != nil { + return nil, fmt.Errorf("revalidate selected OpenAI account: %w", err) + } + if latest == nil || latest.ID != account.ID || + !isOpenAIAccountEligibleForRequest(latest, requirements.RequestedModel, requirements.RequireCompact) || + s.isOpenAIAccountRequestRuntimeBlocked(latest, requirements.RequestedModel) || + !s.isOpenAIAccountTransportCompatible(latest, requirements.RequiredTransport) || + !accountSupportsRequestedOpenAIImageCapability(latest, requirements.RequiredImageCapability) || + (requirements.RequiredEndpointCapability != "" && !latest.SupportsOpenAIEndpointCapability(requirements.RequiredEndpointCapability)) || + (requirements.RequiredPlatform != "" && latest.Platform != requirements.RequiredPlatform) { + return nil, ErrNoAvailableAccounts + } + + if isModeGroup { + return latest, nil + } + if !IsAccountVisibleToRequestUser(ctx, latest) || !s.isOpenAIAccountInRequestGroup(latest, groupID) { + return nil, ErrNoAvailableAccounts + } + return latest, nil +} + func (s *OpenAIGatewayService) getSchedulableAccount(ctx context.Context, accountID int64) (*Account, error) { var ( account *Account @@ -2284,6 +2366,9 @@ func (s *OpenAIGatewayService) hydrateSelectedAccount(ctx context.Context, accou func (s *OpenAIGatewayService) newSelectionResult(ctx context.Context, account *Account, acquired bool, release func(), waitPlan *AccountWaitPlan) (*AccountSelectionResult, error) { hydrated, err := s.hydrateSelectedAccount(ctx, account) if err != nil { + if acquired && release != nil { + release() + } return nil, err } return &AccountSelectionResult{ @@ -2387,6 +2472,9 @@ func (s *OpenAIGatewayService) shouldFailoverUpstreamError(statusCode int) bool } func (s *OpenAIGatewayService) shouldFailoverOpenAIUpstreamResponse(statusCode int, upstreamMsg string, upstreamBody []byte) bool { + if isOpenAIRequestBodyTooLargeError(statusCode, upstreamMsg, upstreamBody) { + return true + } if s.shouldFailoverUpstreamError(statusCode) { return true } @@ -2446,11 +2534,11 @@ func (s *OpenAIGatewayService) handleFailoverSideEffects(ctx context.Context, re func (s *OpenAIGatewayService) handleFailoverSideEffectsForModel(ctx context.Context, resp *http.Response, account *Account, requestedModel string) { body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) - if s.rateLimitService == nil { + if strings.TrimSpace(requestedModel) != "" { + s.handleOpenAIAccountUpstreamErrorForModel(ctx, account, requestedModel, resp.StatusCode, resp.Header, body) return } - if strings.TrimSpace(requestedModel) != "" { - s.rateLimitService.HandleUpstreamErrorForModel(ctx, account, requestedModel, resp.StatusCode, resp.Header, body) + if s.rateLimitService == nil { return } s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, body) @@ -2463,7 +2551,11 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco // ForwardWithAnalysis forwards request to OpenAI API and reuses parsed /responses metadata when available. func (s *OpenAIGatewayService) ForwardWithAnalysis(ctx context.Context, c *gin.Context, account *Account, body []byte, analysis *OpenAIResponsesRequestAnalysis) (*OpenAIForwardResult, error) { + // Keep attempt TTFT separate from the end-to-end first-output budget. The + // scheduler uses FirstTokenMs to learn account health; including local + // queueing there would incorrectly penalize otherwise healthy accounts. startTime := time.Now() + ctx, firstOutputStart := ensureOpenAIFirstOutputStart(ctx) if analysis == nil || !bytes.Equal(analysis.Body, body) { var err error analysis, err = AnalyzeOpenAIResponsesRequest(body) @@ -2884,13 +2976,9 @@ func (s *OpenAIGatewayService) ForwardWithAnalysis(ctx context.Context, c *gin.C if maxOutputTokens, hasMaxOutputTokens := reqBody["max_output_tokens"]; hasMaxOutputTokens { switch account.Platform { case PlatformOpenAI: - // For OpenAI API Key, remove max_output_tokens (not supported) - // For OpenAI OAuth (Responses API), keep it (supported) - if account.Type == AccountTypeAPIKey { - delete(reqBody, "max_output_tokens") - bodyModified = true - markPatchDelete("max_output_tokens") - } + // Responses-native output limits are preserved. Compatible upstreams + // that explicitly reject this field are handled by the bounded retry + // loop below, avoiding unnecessary semantic loss for conforming APIs. case PlatformAnthropic: // For Anthropic (Claude), convert to max_tokens delete(reqBody, "max_output_tokens") @@ -3257,7 +3345,25 @@ func (s *OpenAIGatewayService) ForwardWithAnalysis(ctx context.Context, c *gin.C httpInvalidEncryptedContentRetryTried := false agentIdentityTaskRecoveryTried := agentIdentityTaskRecoveryWasTried(ctx) + rejectedFieldRetryState := newOpenAIResponsesRejectedFieldRetryState(body) for { + // The routing budget includes local queueing and account selection. If it + // is already exhausted, fail before building or dialing another upstream + // request; otherwise a retry would spend network resources after the + // client-visible deadline and incorrectly look like an account timeout. + if firstOutputTimeout > 0 && !time.Now().Before(firstOutputStart.Add(firstOutputTimeout)) { + return nil, s.newOpenAIFirstOutputTimeoutError( + ctx, + c, + account, + firstOutputStart, + originalModel, + reasoningEffortValue, + firstOutputTimeout, + openAIFirstOutputPhaseRoutingBudget, + nil, + ) + } // Build upstream request upstreamCtx, releaseUpstreamCtx := s.detachOpenAIUpstreamContext(ctx) var headerGuard *openAIFirstOutputHeaderGuard @@ -3265,7 +3371,7 @@ func (s *OpenAIGatewayService) ForwardWithAnalysis(ctx context.Context, c *gin.C upstreamCtx, headerGuard = newOpenAIFirstOutputHeaderGuard( upstreamCtx, releaseUpstreamCtx, - startTime.Add(firstOutputTimeout), + firstOutputStart.Add(firstOutputTimeout), ) } upstreamReq, err := s.buildUpstreamRequest(upstreamCtx, c, account, body, token, reqStream, promptCacheKey, isCodexCLI) @@ -3278,6 +3384,22 @@ func (s *OpenAIGatewayService) ForwardWithAnalysis(ctx context.Context, c *gin.C } return nil, err } + if firstOutputTimeout > 0 && !time.Now().Before(firstOutputStart.Add(firstOutputTimeout)) { + if headerGuard != nil { + headerGuard.close() + } + return nil, s.newOpenAIFirstOutputTimeoutError( + ctx, + c, + account, + firstOutputStart, + originalModel, + reasoningEffortValue, + firstOutputTimeout, + openAIFirstOutputPhaseRoutingBudget, + nil, + ) + } expectedAgentIdentityTaskID := strings.TrimSpace(account.GetCredential("task_id")) // Get proxy URL @@ -3300,11 +3422,11 @@ func (s *OpenAIGatewayService) ForwardWithAnalysis(ctx context.Context, c *gin.C ctx, c, account, - startTime, + firstOutputStart, originalModel, reasoningEffortValue, firstOutputTimeout, - "response_headers", + openAIFirstOutputPhaseResponseHeaders, nil, ) } @@ -3347,11 +3469,23 @@ func (s *OpenAIGatewayService) ForwardWithAnalysis(ctx context.Context, c *gin.C } setOpsUpstreamRequestBody(c, body) 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 + if err := json.Unmarshal(body, &openAIReqBody); err != nil { + return nil, fmt.Errorf("decode rejected Responses field retry body: %w", err) + } + setOpsUpstreamRequestBody(c, body) + 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 { @@ -3373,11 +3507,13 @@ func (s *OpenAIGatewayService) ForwardWithAnalysis(ctx context.Context, c *gin.C }) s.handleFailoverSideEffectsForModel(ctx, resp, account, originalModel) - return nil, &UpstreamFailoverError{ - StatusCode: resp.StatusCode, - ResponseBody: respBody, - RetryableOnSameAccount: shouldRetryOpenAIOnSamePoolAccount(account, resp.StatusCode, upstreamMsg, respBody), - } + return nil, newOpenAIUpstreamFailoverError( + resp.StatusCode, + resp.Header, + respBody, + upstreamMsg, + shouldRetryOpenAIOnSamePoolAccount(account, resp.StatusCode, upstreamMsg, respBody), + ) } return s.handleErrorResponse(ctx, resp, c, account, body, originalModel) } @@ -3451,6 +3587,7 @@ func (s *OpenAIGatewayService) forwardOpenAIPassthrough( reqStream bool, startTime time.Time, ) (*OpenAIForwardResult, error) { + ctx, firstOutputStart := ensureOpenAIFirstOutputStart(ctx) cleanRelaySessionBody := body upstreamPassthroughModel := "" if isOpenAIResponsesCompactPath(c) { @@ -3536,6 +3673,14 @@ func (s *OpenAIGatewayService) forwardOpenAIPassthrough( body = updatedBody responseEndpoint := openAIResponsesEndpoint + openAIResponsesRequestPathSuffix(c) imageBillingConfig := resolveOpenAIResponseImageBillingConfigFromBody(responseEndpoint, reqModel, body) + reasoningEffortValue := "" + if reasoningEffort != nil { + reasoningEffortValue = *reasoningEffort + } + firstOutputTimeout := time.Duration(0) + if reqStream && account != nil && account.Platform == PlatformOpenAI { + firstOutputTimeout = s.openAIFirstOutputTimeout(reasoningEffortValue) + } logger.LegacyPrintf("service.openai_gateway", "[OpenAI passthrough] matched passthrough branch: account=%d name=%s type=%s model=%s stream=%v", @@ -3568,12 +3713,54 @@ func (s *OpenAIGatewayService) forwardOpenAIPassthrough( agentIdentityTaskRecoveryTried := agentIdentityTaskRecoveryWasTried(ctx) for { + if firstOutputTimeout > 0 && !time.Now().Before(firstOutputStart.Add(firstOutputTimeout)) { + return nil, s.newOpenAIFirstOutputTimeoutError( + ctx, + c, + account, + firstOutputStart, + reqModel, + reasoningEffortValue, + firstOutputTimeout, + openAIFirstOutputPhaseRoutingBudget, + nil, + ) + } upstreamCtx, releaseUpstreamCtx := s.detachOpenAIUpstreamContext(ctx) + var headerGuard *openAIFirstOutputHeaderGuard + if firstOutputTimeout > 0 { + upstreamCtx, headerGuard = newOpenAIFirstOutputHeaderGuard( + upstreamCtx, + releaseUpstreamCtx, + firstOutputStart.Add(firstOutputTimeout), + ) + } upstreamReq, err := s.buildUpstreamRequestOpenAIPassthrough(upstreamCtx, c, account, body, token) - releaseUpstreamCtx() + if headerGuard == nil { + releaseUpstreamCtx() + } if err != nil { + if headerGuard != nil { + headerGuard.close() + } return nil, err } + if firstOutputTimeout > 0 && !time.Now().Before(firstOutputStart.Add(firstOutputTimeout)) { + if headerGuard != nil { + headerGuard.close() + } + return nil, s.newOpenAIFirstOutputTimeoutError( + ctx, + c, + account, + firstOutputStart, + reqModel, + reasoningEffortValue, + firstOutputTimeout, + openAIFirstOutputPhaseRoutingBudget, + nil, + ) + } expectedAgentIdentityTaskID := strings.TrimSpace(account.GetCredential("task_id")) proxyURL := "" @@ -3590,6 +3777,23 @@ func (s *OpenAIGatewayService) forwardOpenAIPassthrough( upstreamStart := time.Now() resp, err := s.httpUpstream.Do(upstreamReq, proxyURL, account.ID, account.Concurrency) SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds()) + if headerGuard != nil && headerGuard.stopHeaderWait() { + if resp != nil && resp.Body != nil { + _ = resp.Body.Close() + } + headerGuard.close() + return nil, s.newOpenAIFirstOutputTimeoutError( + ctx, + c, + account, + firstOutputStart, + reqModel, + reasoningEffortValue, + firstOutputTimeout, + openAIFirstOutputPhaseResponseHeaders, + nil, + ) + } if err != nil { safeErr := sanitizeUpstreamErrorMessage(err.Error()) setOpsUpstreamError(c, 0, safeErr, "") @@ -3602,6 +3806,9 @@ func (s *OpenAIGatewayService) forwardOpenAIPassthrough( Kind: "request_error", Message: safeErr, }) + if headerGuard != nil { + headerGuard.close() + } c.JSON(http.StatusBadGateway, gin.H{ "error": gin.H{ "type": "upstream_error", @@ -3610,6 +3817,9 @@ func (s *OpenAIGatewayService) forwardOpenAIPassthrough( }) return nil, fmt.Errorf("upstream request failed: %s", safeErr) } + if headerGuard != nil { + resp.Body = &openAIRequestContextReadCloser{ReadCloser: resp.Body, cleanup: headerGuard.close} + } defer func() { _ = resp.Body.Close() }() if resp.StatusCode >= 400 { @@ -3644,7 +3854,7 @@ func (s *OpenAIGatewayService) forwardOpenAIPassthrough( if resp != nil { resp.Request = nil } - result, err := s.handleStreamingResponsePassthrough(ctx, resp, c, account, startTime, reqModel, upstreamPassthroughModel) + result, err := s.handleStreamingResponsePassthroughWithReasoning(ctx, resp, c, account, startTime, reqModel, upstreamPassthroughModel, reasoningEffortValue) if err != nil { return nil, err } @@ -3853,6 +4063,9 @@ func shouldFailoverOpenAIPassthroughResponse(statusCode int, upstreamMsg string, if isOpenAIContextWindowError(upstreamMsg, upstreamBody) { return false } + if isOpenAIRequestBodyTooLargeError(statusCode, upstreamMsg, upstreamBody) { + return true + } switch statusCode { case http.StatusTooManyRequests, 529: return true @@ -3887,7 +4100,7 @@ func (s *OpenAIGatewayService) handleFailoverErrorResponsePassthrough( logOpenAIInstructionsRequiredDebug(ctx, c, account, resp.StatusCode, upstreamMsg, requestBody, body) if s.rateLimitService != nil { requestedModel := extractOpenAIModelFromRequestBody(requestBody) - _ = s.rateLimitService.HandleUpstreamErrorForModel(ctx, account, requestedModel, resp.StatusCode, resp.Header, body) + _ = s.handleOpenAIAccountUpstreamErrorForModel(ctx, account, requestedModel, resp.StatusCode, resp.Header, body) } appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ Platform: account.Platform, @@ -3901,11 +4114,7 @@ func (s *OpenAIGatewayService) handleFailoverErrorResponsePassthrough( Detail: upstreamDetail, UpstreamResponseBody: upstreamDetail, }) - return &UpstreamFailoverError{ - StatusCode: resp.StatusCode, - ResponseBody: body, - ResponseHeaders: resp.Header.Clone(), - } + return newOpenAIUpstreamFailoverError(resp.StatusCode, resp.Header, body, upstreamMsg, false) } func (s *OpenAIGatewayService) handleErrorResponsePassthrough( @@ -3944,7 +4153,7 @@ func (s *OpenAIGatewayService) handleErrorResponsePassthrough( // account state still needs to be updated so sticky routing can stop // reusing a freshly rate-limited account. requestedModel := extractOpenAIModelFromRequestBody(requestBody) - _ = s.rateLimitService.HandleUpstreamErrorForModel(ctx, account, requestedModel, resp.StatusCode, resp.Header, body) + _ = s.handleOpenAIAccountUpstreamErrorForModel(ctx, account, requestedModel, resp.StatusCode, resp.Header, body) } appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ Platform: account.Platform, @@ -4019,6 +4228,20 @@ type openaiStreamingResultPassthrough struct { responseServiceTier string } +// openAIStreamEventDisposition keeps protocol delivery separate from the +// first visible/meaningful output signal. A lifecycle event may need to be +// eligible for protocol delivery (for example a tool name), while it must not +// stop the first-output watchdog or be reported as a first token. Guarded +// failover keeps protocol-only events staged until semantic/terminal output. +// Conversely, a +// terminal event with an empty output must stop the watchdog and flush the +// staged protocol bytes without fabricating a FirstTokenMs value. +type openAIStreamEventDisposition struct { + commitPending bool + semanticOutput bool + timerSatisfied bool +} + func openAIStreamClientOutputStarted(c *gin.Context, localStarted bool) bool { if localStarted { return true @@ -4026,24 +4249,214 @@ func openAIStreamClientOutputStarted(c *gin.Context, localStarted bool) bool { return c != nil && c.Writer != nil && c.Writer.Written() } -func openAIStreamEventIsPreamble(eventType string) bool { - switch strings.TrimSpace(eventType) { - case "response.created", "response.in_progress": - return true +func openAIStreamStringFieldHasValue(data string, paths ...string) bool { + for _, path := range paths { + value := gjson.Get(data, path) + if value.Exists() && value.Type == gjson.String && value.String() != "" { + return true + } + } + return false +} + +func openAIStreamOpaqueFieldHasValue(data string, paths ...string) bool { + for _, path := range paths { + value := gjson.Get(data, path) + if value.Exists() && value.Type == gjson.String && strings.TrimSpace(value.String()) != "" { + return true + } + } + return false +} + +func openAIStreamStructuredFieldHasValue(data string, paths ...string) bool { + for _, path := range paths { + value := gjson.Get(data, path) + if !value.Exists() { + continue + } + switch value.Type { + case gjson.String: + if value.String() != "" { + return true + } + case gjson.JSON: + raw := strings.TrimSpace(value.Raw) + if raw != "" && raw != "null" && raw != "{}" && raw != "[]" { + return true + } + case gjson.Number, gjson.True, gjson.False: + return true + } + } + return false +} + +func openAIStreamToolCallsHaveMetadata(item gjson.Result) bool { + for _, call := range item.Get("tool_calls").Array() { + if openAIStreamStringFieldHasValue(call.Raw, "name", "function.name") || + openAIStreamStructuredFieldHasValue(call.Raw, "arguments", "function.arguments") { + return true + } + } + return false +} + +func openAIStreamOutputItemHasContent(item gjson.Result) (semantic, protocol bool) { + if !item.Exists() || !item.IsObject() { + return false, false + } + itemType := strings.ToLower(strings.TrimSpace(item.Get("type").String())) + switch itemType { + case "message": + for _, part := range item.Get("content").Array() { + if openAIStreamStringFieldHasValue(part.Raw, "text", "refusal", "transcript") || + openAIStreamOpaqueFieldHasValue(part.Raw, "audio", "data", "file_id", "image_url") { + return true, true + } + } + return openAIStreamStringFieldHasValue(item.Raw, "text", "refusal"), false + case "function_call", "custom_tool_call", "tool_call": + if openAIStreamStructuredFieldHasValue(item.Raw, "arguments", "input", "function.arguments", "custom_tool_call.input") { + return true, true + } + // A tool name is protocol-significant, but it is not a generated token by + // itself. Guarded streams keep it staged until failover is no longer safe. + return false, openAIStreamStringFieldHasValue(item.Raw, "name", "function.name") || openAIStreamToolCallsHaveMetadata(item) + case "image_generation_call": + semantic = openAIStreamOpaqueFieldHasValue(item.Raw, "result", "partial_image_b64", "image", "data") + return semantic, semantic + case "tool_search_call", "mcp_call", "computer_call", "local_shell_call", "shell_call", "web_search_call", "file_search_call", "code_interpreter_call": + semantic = openAIStreamStructuredFieldHasValue(item.Raw, "arguments", "input", "action", "code", "output", "result") + protocol = semantic || openAIStreamStringFieldHasValue(item.Raw, "name", "call_id", "server_label") + return semantic, protocol + case "compaction", "compaction_summary": + semantic = openAIStreamStringFieldHasValue(item.Raw, "summary", "text", "content", "encrypted_content") + return semantic, semantic default: - return false + semantic = openAIStreamStringFieldHasValue(item.Raw, "text", "refusal", "transcript") || + openAIStreamStructuredFieldHasValue(item.Raw, "arguments", "input", "output") || + openAIStreamOpaqueFieldHasValue(item.Raw, "audio", "partial_image_b64", "result") + if !semantic { + semantic = openAIStreamStringFieldHasValue(item.Raw, "summary", "content") + } + return semantic, semantic } } -func openAIStreamDataStartsClientOutput(data, eventType string) bool { +func openAIStreamResponseHasContent(data string) bool { + for _, path := range []string{"response.output", "output"} { + output := gjson.Get(data, path) + if !output.Exists() { + continue + } + if output.IsArray() { + for _, item := range output.Array() { + semantic, _ := openAIStreamOutputItemHasContent(item) + if semantic { + return true + } + } + } + } + return openAIStreamStringFieldHasValue(data, "response.output_text", "response.refusal", "response.text", "response.audio") +} + +func openAIStreamEventPayloadHasContent(data, eventType string) (semantic, protocol bool) { + switch eventType { + case "response.output_item.added", "response.output_item.done": + semantic, protocol = openAIStreamOutputItemHasContent(gjson.Get(data, "item")) + // A completed function/custom-tool item with only a name is still a + // meaningful model decision (some probe/tool providers omit arguments). + // The corresponding `added` lifecycle event remains protocol-only. + if eventType == "response.output_item.done" && !semantic && protocol { + itemType := strings.ToLower(strings.TrimSpace(gjson.Get(data, "item.type").String())) + if itemType == "function_call" || itemType == "custom_tool_call" || itemType == "tool_call" { + semantic = true + } + } + // The lifecycle event itself is protocol-significant even when the item + // is only an empty message shell. Unguarded streams must preserve it; + // guarded streams still stage it because it is not semantic output. + return semantic, true + case "response.content_part.added", "response.content_part.done": + part := gjson.Get(data, "part") + semantic = openAIStreamStringFieldHasValue(part.Raw, "text", "refusal", "transcript") || + openAIStreamOpaqueFieldHasValue(part.Raw, "audio", "data", "file_id", "image_url") + return semantic, true + case "response.output_text.delta", "response.refusal.delta", "response.function_call_arguments.delta", "response.custom_tool_call_input.delta", "response.reasoning_summary_text.delta": + semantic = openAIStreamStringFieldHasValue(data, "delta") + return semantic, semantic + case "response.output_audio.delta": + semantic = openAIStreamOpaqueFieldHasValue(data, "delta") + return semantic, semantic + case "response.output_text.done": + semantic = openAIStreamStringFieldHasValue(data, "text") + return semantic, semantic + case "response.refusal.done": + semantic = openAIStreamStringFieldHasValue(data, "refusal") + return semantic, semantic + case "response.function_call_arguments.done": + semantic = openAIStreamStringFieldHasValue(data, "arguments") + return semantic, semantic + case "response.custom_tool_call_input.done": + semantic = openAIStreamStringFieldHasValue(data, "input") + return semantic, semantic + case "response.reasoning_summary_part.done": + semantic = openAIStreamStringFieldHasValue(data, "part.text", "text") + return semantic, semantic + case "response.image_generation_call.partial_image": + semantic = openAIStreamOpaqueFieldHasValue(data, "partial_image_b64") + return semantic, semantic + case "response.image_generation_call.completed": + semantic = openAIStreamOpaqueFieldHasValue(data, "result", "partial_image_b64", "image", "data") + return semantic, semantic + } + + // Keep unknown future lifecycle events fail-closed. A known delta/done + // event can still be recognized by its payload without treating arbitrary + // metadata as a first token. + if strings.HasSuffix(eventType, ".delta") { + semantic = openAIStreamStringFieldHasValue(data, "delta") || openAIStreamOpaqueFieldHasValue(data, "audio", "partial_image_b64") + } else if strings.HasSuffix(eventType, ".done") { + semantic = openAIStreamStringFieldHasValue(data, "text", "refusal", "transcript", "code") || + openAIStreamStructuredFieldHasValue(data, "arguments", "input") || + openAIStreamOpaqueFieldHasValue(data, "audio", "result", "partial_image_b64") + } + return semantic, semantic +} + +func classifyOpenAIStreamEvent(data, eventType string) openAIStreamEventDisposition { trimmed := strings.TrimSpace(data) if trimmed == "" { - return false - } - if strings.TrimSpace(eventType) == "response.failed" { - return false + return openAIStreamEventDisposition{} + } + if trimmed == "[DONE]" { + return openAIStreamEventDisposition{commitPending: true, timerSatisfied: true} + } + if !gjson.Valid(trimmed) { + return openAIStreamEventDisposition{} + } + eventType = strings.TrimSpace(eventType) + if eventType == "" { + eventType = strings.TrimSpace(gjson.Get(trimmed, "type").String()) + } + disposition := openAIStreamEventDisposition{} + switch eventType { + case "response.created", "response.in_progress", "response.reasoning_summary_part.added": + return disposition + case "response.failed", "error", "response.error": + return openAIStreamEventDisposition{commitPending: true, timerSatisfied: true} + case "response.completed", "response.done", "response.incomplete", "response.cancelled", "response.canceled": + disposition.semanticOutput = openAIStreamResponseHasContent(trimmed) + disposition.commitPending = true + disposition.timerSatisfied = true + return disposition + default: + disposition.semanticOutput, disposition.commitPending = openAIStreamEventPayloadHasContent(trimmed, eventType) + disposition.timerSatisfied = disposition.semanticOutput + return disposition } - return !openAIStreamEventIsPreamble(eventType) } func openAIStreamFailedEventShouldFailover(payload []byte, message string) bool { @@ -4178,16 +4591,57 @@ func (s *OpenAIGatewayService) handleStreamingResponsePassthrough( originalModel string, mappedModel string, ) (*openaiStreamingResultPassthrough, error) { - writeOpenAIPassthroughResponseHeaders(c.Writer.Header(), resp.Header, s.responseHeaderFilter) + return s.handleStreamingResponsePassthroughWithReasoning(ctx, resp, c, account, startTime, originalModel, mappedModel, "") +} + +func (s *OpenAIGatewayService) handleStreamingResponsePassthroughWithReasoning( + ctx context.Context, + resp *http.Response, + c *gin.Context, + account *Account, + startTime time.Time, + originalModel string, + mappedModel string, + reasoningEffort string, +) (*openaiStreamingResultPassthrough, error) { + ctx, firstOutputStart := ensureOpenAIFirstOutputStart(ctx) + firstOutputTimeout := time.Duration(0) + if account != nil && account.Platform == PlatformOpenAI { + firstOutputTimeout = s.openAIFirstOutputTimeout(reasoningEffort) + } + guardFirstOutput := firstOutputTimeout > 0 + + attemptResponseHeaders := c.Writer.Header() + if guardFirstOutput { + attemptResponseHeaders = make(http.Header) + } + writeOpenAIPassthroughResponseHeaders(attemptResponseHeaders, resp.Header, s.responseHeaderFilter) + if v := resp.Header.Get("x-request-id"); v != "" { + attemptResponseHeaders.Set("x-request-id", v) + } + applyAttemptResponseHeaders := func() { + if !guardFirstOutput || attemptResponseHeaders == nil { + return + } + dst := c.Writer.Header() + for key, values := range attemptResponseHeaders { + dst.Del(key) + for _, value := range values { + dst.Add(key, value) + } + } + c.Header("Content-Type", "text/event-stream") + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + c.Header("X-Accel-Buffering", "no") + attemptResponseHeaders = nil + } // SSE headers c.Header("Content-Type", "text/event-stream") c.Header("Cache-Control", "no-cache") c.Header("Connection", "keep-alive") c.Header("X-Accel-Buffering", "no") - if v := resp.Header.Get("x-request-id"); v != "" { - c.Header("x-request-id", v) - } w := c.Writer flusher, ok := w.(http.Flusher) @@ -4205,8 +4659,19 @@ func (s *OpenAIGatewayService) handleStreamingResponsePassthrough( failedMessage := "" clientOutputStarted := false upstreamRequestID := strings.TrimSpace(resp.Header.Get("x-request-id")) + var disconnectedDrainStarted bool + cancelDisconnectedDrain := func() {} + defer func() { cancelDisconnectedDrain() }() + startDisconnectedDrain := func() { + if disconnectedDrainStarted { + return + } + disconnectedDrainStarted = true + cancelDisconnectedDrain = s.startDisconnectedStreamDrainDeadline(ctx, resp.Body, upstreamRequestID) + } imageCounter := newOpenAIImageOutputCounter() pendingLines := make([]string, 0, 8) + var pendingLineBytes int64 flushPending := false flushPendingOutput := func() { if clientDisconnected || !flushPending { @@ -4224,14 +4689,53 @@ func (s *OpenAIGatewayService) handleStreamingResponsePassthrough( for _, pending := range pendingLines { if _, err := fmt.Fprintln(w, pending); err != nil { clientDisconnected = true + startDisconnectedDrain() s.stopUpstreamOnClientDisconnect(ctx, resp.Body) s.legacyLogClientDisconnectDrainDecision(ctx, "[OpenAI passthrough] Client disconnected during streaming, continue draining upstream for usage: account=%d", account.ID) return false } } pendingLines = pendingLines[:0] + pendingLineBytes = 0 return true } + appendPendingLine := func(line string) error { + incoming := int64(len(line) + 1) + if incoming > openAIFirstOutputStageMaxBytes-pendingLineBytes { + return fmt.Errorf("%w: buffered=%d incoming=%d limit=%d", errOpenAIFirstOutputStageLimit, pendingLineBytes, incoming, openAIFirstOutputStageMaxBytes) + } + pendingLines = append(pendingLines, line) + pendingLineBytes += incoming + return nil + } + + var firstOutputScanGuard atomic.Bool + firstOutputScanGuard.Store(guardFirstOutput) + var firstOutputTimer *time.Timer + var firstOutputCh <-chan time.Time + if firstOutputTimeout > 0 { + remaining := time.Until(firstOutputStart.Add(firstOutputTimeout)) + if remaining <= 0 { + remaining = time.Nanosecond + } + firstOutputTimer = time.NewTimer(remaining) + firstOutputCh = firstOutputTimer.C + defer firstOutputTimer.Stop() + } + stopFirstOutputTimer := func() { + firstOutputScanGuard.Store(false) + if firstOutputTimer == nil { + return + } + if !firstOutputTimer.Stop() { + select { + case <-firstOutputTimer.C: + default: + } + } + firstOutputTimer = nil + firstOutputCh = nil + } scanner := bufio.NewScanner(resp.Body) maxLineSize := defaultMaxLineSize @@ -4240,13 +4744,142 @@ func (s *OpenAIGatewayService) handleStreamingResponsePassthrough( } scanBuf := getSSEScannerBuf64K() scanner.Buffer(scanBuf[:0], maxLineSize) - defer putSSEScannerBuf64K(scanBuf) + if guardFirstOutput { + scanner.Split(openAIFirstOutputDynamicScanLines(&firstOutputScanGuard)) + } documentScanner := newOpenAISSEJSONDocumentScanner(scanner) + streamInterval := time.Duration(0) + if s.cfg != nil && s.cfg.Gateway.StreamDataIntervalTimeout > 0 { + streamInterval = time.Duration(s.cfg.Gateway.StreamDataIntervalTimeout) * time.Second + } + var streamIdleTimer *time.Timer + var streamIdleCh <-chan time.Time + if streamInterval > 0 { + streamIdleTimer = time.NewTimer(streamInterval) + streamIdleCh = streamIdleTimer.C + defer streamIdleTimer.Stop() + } + var lastUpstreamReadAt atomic.Int64 + lastUpstreamReadAt.Store(time.Now().UnixNano()) + type passthroughScanEvent struct { + line string + err error + } + scanEvents := make(chan passthroughScanEvent, 1) + stopScan := make(chan struct{}) + go func(scanBuf *sseScannerBuf64K) { + defer putSSEScannerBuf64K(scanBuf) + defer close(scanEvents) + send := func(event passthroughScanEvent) bool { + select { + case scanEvents <- event: + return true + case <-stopScan: + return false + } + } + for documentScanner.Scan() { + lastUpstreamReadAt.Store(time.Now().UnixNano()) + if !send(passthroughScanEvent{line: documentScanner.Text()}) { + return + } + } + if err := documentScanner.Err(); err != nil { + _ = send(passthroughScanEvent{err: err}) + } + }(scanBuf) + defer close(stopScan) needModelReplace := strings.TrimSpace(originalModel) != "" && strings.TrimSpace(mappedModel) != "" && strings.TrimSpace(originalModel) != strings.TrimSpace(mappedModel) + guardedEventInProgress := false + guardedEventHasSemanticOutput := false + guardedEventSatisfiesFirstOutput := false + completeGuardedEvent := func() { + semanticOutput := guardedEventHasSemanticOutput + satisfiesFirstOutput := guardedEventSatisfiesFirstOutput + guardedEventInProgress = false + guardedEventHasSemanticOutput = false + guardedEventSatisfiesFirstOutput = false + if !semanticOutput && !satisfiesFirstOutput { + return + } + applyAttemptResponseHeaders() + if !writePendingLines() { + stopFirstOutputTimer() + guardFirstOutput = false + return + } + clientOutputStarted = true + flusher.Flush() + flushPending = false + if semanticOutput && firstTokenMs == nil { + ms := int(time.Since(startTime).Milliseconds()) + firstTokenMs = &ms + } + stopFirstOutputTimer() + guardFirstOutput = false + } - for documentScanner.Scan() { - line := documentScanner.Text() + var scanErr error +streamLoop: + for { + var line string + select { + case event, ok := <-scanEvents: + if !ok { + break streamLoop + } + if event.err != nil { + scanErr = event.err + break streamLoop + } + line = event.line + case <-firstOutputCh: + _ = resp.Body.Close() + return resultWithUsage(), s.newOpenAIFirstOutputTimeoutError( + ctx, + c, + account, + firstOutputStart, + originalModel, + reasoningEffort, + firstOutputTimeout, + openAIFirstOutputPhaseSemanticOutput, + resp.Header, + ) + case <-streamIdleCh: + lastRead := time.Unix(0, lastUpstreamReadAt.Load()) + idleFor := time.Since(lastRead) + if idleFor < streamInterval { + streamIdleTimer.Reset(streamInterval - idleFor) + continue + } + _ = resp.Body.Close() + if clientDisconnected { + return resultWithUsage(), errors.New("stream usage incomplete after timeout") + } + if sawTerminalEvent && !sawFailedEvent { + return resultWithUsage(), nil + } + logger.LegacyPrintf("service.openai_gateway", "[OpenAI passthrough] Stream data interval timeout: account=%d model=%s interval=%s", account.ID, originalModel, streamInterval) + if guardFirstOutput && !clientOutputStarted { + failoverErr := s.newOpenAIStreamFailoverError(c, account, true, upstreamRequestID, nil, "OpenAI stream produced no complete output event before the data interval timeout") + failoverErr.SafeToFailoverAfterWrite = true + return resultWithUsage(), failoverErr + } + if s.rateLimitService != nil { + s.rateLimitService.HandleStreamTimeout(ctx, account, originalModel) + } + payload := `{"type":"error","sequence_number":0,"error":{"type":"upstream_error","message":"stream_timeout","code":"stream_timeout"}}` + if _, err := fmt.Fprintln(w, "data: "+payload); err == nil { + _, _ = fmt.Fprintln(w) + flusher.Flush() + } else { + clientDisconnected = true + startDisconnectedDrain() + } + return resultWithUsage(), errors.New("stream data interval timeout") + } lineStartsClientOutput := false forceFlushFailedEvent := false if data, ok := extractOpenAISSEDataLine(line); ok { @@ -4276,7 +4909,7 @@ func (s *OpenAIGatewayService) handleStreamingResponsePassthrough( } } eventType := strings.TrimSpace(gjson.Get(trimmedData, "type").String()) - if eventType == "response.failed" { + if eventType == "response.failed" || eventType == "error" || eventType == "response.error" { failedMessage = extractOpenAISSEErrorMessage(dataBytes) s.parseSSEUsageBytes(dataBytes, usage) if hit, code, msg := detectOpenAICyberPolicy(dataBytes); hit { @@ -4289,13 +4922,17 @@ func (s *OpenAIGatewayService) handleStreamingResponsePassthrough( UpstreamOutTok: usage.OutputTokens, }) } - if !openAIStreamClientOutputStarted(c, clientOutputStarted) { - if status, errType, errMsg, matched := applyOpenAIStreamFailedErrorPassthroughRule(c, account.Platform, dataBytes, failedMessage); matched { - s.recordOpenAIStreamUpstreamError(c, account, true, upstreamRequestID, "http_error", dataBytes, failedMessage) - MarkResponseCommitted(c) - c.Writer.Header().Set("Content-Type", "application/json; charset=utf-8") - c.JSON(status, gin.H{"error": gin.H{"type": errType, "message": errMsg}}) - return resultWithUsage(), fmt.Errorf("upstream response failed: passthrough rule matched message=%s", errMsg) + if !clientOutputStarted { + // Routing/user-slot pings are neutral and do not make an upstream + // attempt unsafe to replay. Avoid JSON only when bytes already exist. + if !openAIStreamClientOutputStarted(c, false) { + if status, errType, errMsg, matched := applyOpenAIStreamFailedErrorPassthroughRule(c, account.Platform, dataBytes, failedMessage); matched { + s.recordOpenAIStreamUpstreamError(c, account, true, upstreamRequestID, "http_error", dataBytes, failedMessage) + MarkResponseCommitted(c) + c.Writer.Header().Set("Content-Type", "application/json; charset=utf-8") + c.JSON(status, gin.H{"error": gin.H{"type": errType, "message": errMsg}}) + return resultWithUsage(), fmt.Errorf("upstream response failed: passthrough rule matched message=%s", errMsg) + } } if openAIStreamFailedEventShouldFailover(dataBytes, failedMessage) { return resultWithUsage(), s.newOpenAIStreamFailoverError(c, account, true, upstreamRequestID, dataBytes, failedMessage) @@ -4325,17 +4962,40 @@ func (s *OpenAIGatewayService) handleStreamingResponsePassthrough( trimmedData = strings.TrimSpace(string(sanitizedData)) line = "data: " + string(sanitizedData) } - lineStartsClientOutput = forceFlushFailedEvent || openAIStreamDataStartsClientOutput(trimmedData, eventType) - if firstTokenMs == nil && lineStartsClientOutput && trimmedData != "[DONE]" { + disposition := classifyOpenAIStreamEvent(trimmedData, eventType) + lineStartsClientOutput = forceFlushFailedEvent || disposition.commitPending + if guardFirstOutput { + guardedEventHasSemanticOutput = guardedEventHasSemanticOutput || disposition.semanticOutput + guardedEventSatisfiesFirstOutput = guardedEventSatisfiesFirstOutput || forceFlushFailedEvent || disposition.timerSatisfied + lineStartsClientOutput = false + } else if firstTokenMs == nil && disposition.semanticOutput { ms := int(time.Since(startTime).Milliseconds()) firstTokenMs = &ms + stopFirstOutputTimer() } s.parseSSEUsageBytes(dataBytes, usage) } if !clientDisconnected { + if guardFirstOutput { + if err := appendPendingLine(line); err != nil { + failoverErr := s.newOpenAIStreamFailoverError(c, account, true, upstreamRequestID, nil, "OpenAI passthrough first-output staging limit exceeded") + failoverErr.SafeToFailoverAfterWrite = true + return resultWithUsage(), failoverErr + } + if line == "" { + completeGuardedEvent() + } else { + guardedEventInProgress = true + } + continue + } if !clientOutputStarted && !lineStartsClientOutput { - pendingLines = append(pendingLines, line) + if err := appendPendingLine(line); err != nil { + failoverErr := s.newOpenAIStreamFailoverError(c, account, true, upstreamRequestID, nil, "OpenAI passthrough first-output staging limit exceeded") + failoverErr.SafeToFailoverAfterWrite = true + return resultWithUsage(), failoverErr + } continue } if !clientOutputStarted && len(pendingLines) > 0 { @@ -4345,6 +5005,7 @@ func (s *OpenAIGatewayService) handleStreamingResponsePassthrough( } if _, err := fmt.Fprintln(w, line); err != nil { clientDisconnected = true + startDisconnectedDrain() s.stopUpstreamOnClientDisconnect(ctx, resp.Body) s.legacyLogClientDisconnectDrainDecision(ctx, "[OpenAI passthrough] Client disconnected during streaming, continue draining upstream for usage: account=%d", account.ID) } else { @@ -4356,7 +5017,12 @@ func (s *OpenAIGatewayService) handleStreamingResponsePassthrough( } } } - if err := documentScanner.Err(); err != nil { + if scanErr == nil && guardFirstOutput && guardedEventInProgress { + // EOF dispatches the final complete SSE event even when the upstream + // omits the trailing blank line. + completeGuardedEvent() + } + if err := scanErr; err != nil { if clientDisconnected { streamErr := s.clientDisconnectIncompleteUsageError(ctx) if streamErr == nil { @@ -4377,11 +5043,17 @@ func (s *OpenAIGatewayService) handleStreamingResponsePassthrough( if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { return resultWithUsage(), fmt.Errorf("stream usage incomplete: %w", err) } + if errors.Is(err, errOpenAIFirstOutputScannerLimit) && firstTokenMs == nil { + logger.LegacyPrintf("service.openai_gateway", "[OpenAI passthrough] SSE token exceeded guarded first-output limit: account=%d limit=%d error=%v", account.ID, openAIFirstOutputStageMaxBytes+openAIFirstOutputScannerFramingAllowance, err) + failoverErr := s.newOpenAIStreamFailoverError(c, account, true, upstreamRequestID, nil, "OpenAI SSE line exceeds guarded first-output limit") + failoverErr.SafeToFailoverAfterWrite = true + return resultWithUsage(), failoverErr + } if errors.Is(err, bufio.ErrTooLong) { logger.LegacyPrintf("service.openai_gateway", "[OpenAI passthrough] SSE line too long: account=%d max_size=%d error=%v", account.ID, maxLineSize, err) return resultWithUsage(), err } - if !openAIStreamClientOutputStarted(c, clientOutputStarted) { + if !clientOutputStarted { msg := "OpenAI stream disconnected before completion" if errText := strings.TrimSpace(err.Error()); errText != "" { msg += ": " + errText @@ -4406,7 +5078,7 @@ func (s *OpenAIGatewayService) handleStreamingResponsePassthrough( zap.Int64("account_id", account.ID), zap.String("upstream_request_id", upstreamRequestID), ).Info("OpenAI passthrough upstream stream ended before [DONE], suspected truncated stream") - if !openAIStreamClientOutputStarted(c, clientOutputStarted) { + if !clientOutputStarted { return resultWithUsage(), s.newOpenAIStreamFailoverError(c, account, true, upstreamRequestID, nil, "OpenAI stream ended before a terminal event") } @@ -4842,7 +5514,7 @@ func (s *OpenAIGatewayService) handleErrorResponse( // Handle upstream error (mark account status) shouldDisable := false if s.rateLimitService != nil { - shouldDisable = s.rateLimitService.HandleUpstreamErrorForModel(ctx, account, requestedModel, resp.StatusCode, resp.Header, body) + shouldDisable = s.handleOpenAIAccountUpstreamErrorForModel(ctx, account, requestedModel, resp.StatusCode, resp.Header, body) } kind := "http_error" if shouldDisable { @@ -4985,7 +5657,7 @@ func (s *OpenAIGatewayService) handleCompatErrorResponse( // Track rate limits and decide whether to trigger secondary failover. shouldDisable := false if s.rateLimitService != nil { - shouldDisable = s.rateLimitService.HandleUpstreamErrorForModel( + shouldDisable = s.handleOpenAIAccountUpstreamErrorForModel( c.Request.Context(), account, requestedModel, resp.StatusCode, resp.Header, body, ) } @@ -5044,6 +5716,7 @@ func (s *OpenAIGatewayService) handleStreamingResponseWithReasoning(ctx context. if account != nil && account.Platform == PlatformOpenAI { firstOutputTimeout = s.openAIFirstOutputTimeout(reasoningEffort) } + ctx, firstOutputStart := ensureOpenAIFirstOutputStart(ctx) guardFirstOutput := firstOutputTimeout > 0 var attemptResponseHeaders http.Header if guardFirstOutput { @@ -5173,7 +5846,7 @@ func (s *OpenAIGatewayService) handleStreamingResponseWithReasoning(ctx context. var firstOutputTimer *time.Timer var firstOutputCh <-chan time.Time if firstOutputTimeout > 0 { - remaining := time.Until(startTime.Add(firstOutputTimeout)) + remaining := time.Until(firstOutputStart.Add(firstOutputTimeout)) if remaining <= 0 { remaining = time.Nanosecond } @@ -5207,6 +5880,7 @@ func (s *OpenAIGatewayService) handleStreamingResponseWithReasoning(ctx context. sawFailedEvent := false failedMessage := "" clientOutputStarted := false + neutralKeepaliveWritten := false upstreamRequestID := strings.TrimSpace(resp.Header.Get("x-request-id")) var cancelDisconnectedDrain context.CancelFunc defer func() { @@ -5221,7 +5895,8 @@ func (s *OpenAIGatewayService) handleStreamingResponseWithReasoning(ctx context. } var streamFailoverErr error eventInProgress := false - eventStartsClientOutput := false + eventHasSemanticOutput := false + eventSatisfiesFirstOutput := false eventShouldFlush := false handlePendingWriteError := func(err error) { if firstOutputStage != nil && firstTokenMs == nil && !firstOutputStage.closed { @@ -5242,11 +5917,12 @@ func (s *OpenAIGatewayService) handleStreamingResponseWithReasoning(ctx context. s.legacyLogClientDisconnectDrainDecision(ctx, "Client disconnected during streaming, continuing to drain upstream for billing") } completeGuardedEvent := func(queueDrained bool) { - completedSemanticEvent := eventStartsClientOutput - shouldFlush := eventShouldFlush || (queueDrained && clientOutputStarted) + completedSemanticEvent := eventHasSemanticOutput + commitGuardedEvent := eventHasSemanticOutput || eventSatisfiesFirstOutput + shouldFlush := eventShouldFlush || commitGuardedEvent || (queueDrained && clientOutputStarted) eventInProgress = false if !clientDisconnected { - if completedSemanticEvent { + if commitGuardedEvent { applyAttemptResponseHeaders() } if shouldFlush { @@ -5267,7 +5943,12 @@ func (s *OpenAIGatewayService) handleStreamingResponseWithReasoning(ctx context. firstTokenMs = &ms stopFirstOutputTimer() } - eventStartsClientOutput = false + if eventSatisfiesFirstOutput { + firstOutputScanGuard.Store(false) + stopFirstOutputTimer() + } + eventHasSemanticOutput = false + eventSatisfiesFirstOutput = false eventShouldFlush = false } sendErrorEvent := func(reason string) { @@ -5329,8 +6010,8 @@ func (s *OpenAIGatewayService) handleStreamingResponseWithReasoning(ctx context. return resultWithUsage(), streamErr } } - if !sawTerminalEvent && !openAIStreamClientOutputStarted(c, clientOutputStarted) && !eventShouldFlush { - return resultWithUsage(), s.newOpenAIStreamFailoverError( + if !sawTerminalEvent && !clientOutputStarted && !eventShouldFlush { + failoverErr := s.newOpenAIStreamFailoverError( c, account, false, @@ -5338,6 +6019,8 @@ func (s *OpenAIGatewayService) handleStreamingResponseWithReasoning(ctx context. nil, "OpenAI stream ended before a terminal event", ) + failoverErr.SafeToFailoverAfterWrite = neutralKeepaliveWritten + return resultWithUsage(), failoverErr } flushPending("Client disconnected during final flush, returning collected usage") if !sawTerminalEvent { @@ -5397,12 +6080,14 @@ func (s *OpenAIGatewayService) handleStreamingResponseWithReasoning(ctx context. sendErrorEvent("response_too_large") return resultWithUsage(), scanErr, true } - if !openAIStreamClientOutputStarted(c, clientOutputStarted) && !eventShouldFlush { + if !clientOutputStarted && !eventShouldFlush { msg := "OpenAI stream disconnected before completion" if errText := strings.TrimSpace(scanErr.Error()); errText != "" { msg += ": " + errText } - return resultWithUsage(), s.newOpenAIStreamFailoverError(c, account, false, upstreamRequestID, nil, msg), true + failoverErr := s.newOpenAIStreamFailoverError(c, account, false, upstreamRequestID, nil, msg) + failoverErr.SafeToFailoverAfterWrite = neutralKeepaliveWritten + return resultWithUsage(), failoverErr, true } sendErrorEvent("stream_read_error") return resultWithUsage(), fmt.Errorf("stream read error: %w", scanErr), true @@ -5431,7 +6116,7 @@ func (s *OpenAIGatewayService) handleStreamingResponseWithReasoning(ctx context. } eventType := strings.TrimSpace(gjson.GetBytes(dataBytes, "type").String()) forceFlushFailedEvent := false - if eventType == "response.failed" { + if eventType == "response.failed" || eventType == "error" || eventType == "response.error" { failedMessage = extractOpenAISSEErrorMessage(dataBytes) s.parseSSEUsageBytes(dataBytes, usage) if hit, code, msg := detectOpenAICyberPolicy(dataBytes); hit { @@ -5444,19 +6129,25 @@ func (s *OpenAIGatewayService) handleStreamingResponseWithReasoning(ctx context. UpstreamOutTok: usage.OutputTokens, }) } - if !openAIStreamClientOutputStarted(c, clientOutputStarted) { - if status, errType, errMsg, matched := applyOpenAIStreamFailedErrorPassthroughRule(c, account.Platform, dataBytes, failedMessage); matched { - sawFailedEvent = true - s.recordOpenAIStreamUpstreamError(c, account, false, upstreamRequestID, "http_error", dataBytes, failedMessage) - MarkResponseCommitted(c) - c.Writer.Header().Set("Content-Type", "application/json; charset=utf-8") - c.JSON(status, gin.H{"error": gin.H{"type": errType, "message": errMsg}}) - streamFailoverErr = fmt.Errorf("upstream response failed: passthrough rule matched message=%s", errMsg) - return + if !clientOutputStarted { + // A neutral keepalive is not an upstream response event. Preserve + // safe failover, but do not write a JSON error after SSE bytes. + if !openAIStreamClientOutputStarted(c, false) { + if status, errType, errMsg, matched := applyOpenAIStreamFailedErrorPassthroughRule(c, account.Platform, dataBytes, failedMessage); matched { + sawFailedEvent = true + s.recordOpenAIStreamUpstreamError(c, account, false, upstreamRequestID, "http_error", dataBytes, failedMessage) + MarkResponseCommitted(c) + c.Writer.Header().Set("Content-Type", "application/json; charset=utf-8") + c.JSON(status, gin.H{"error": gin.H{"type": errType, "message": errMsg}}) + streamFailoverErr = fmt.Errorf("upstream response failed: passthrough rule matched message=%s", errMsg) + return + } } if openAIStreamFailedEventShouldFailover(dataBytes, failedMessage) { sawFailedEvent = true - streamFailoverErr = s.newOpenAIStreamFailoverError(c, account, false, upstreamRequestID, dataBytes, failedMessage) + failoverErr := s.newOpenAIStreamFailoverError(c, account, false, upstreamRequestID, dataBytes, failedMessage) + failoverErr.SafeToFailoverAfterWrite = neutralKeepaliveWritten + streamFailoverErr = failoverErr return } } else { @@ -5507,15 +6198,22 @@ func (s *OpenAIGatewayService) handleStreamingResponseWithReasoning(ctx context. if needModelReplace && mappedModel != "" && strings.Contains(line, mappedModel) { line = s.replaceModelInSSELine(line, mappedModel, originalModel) } - startsClientOutput := forceFlushFailedEvent || openAIStreamDataStartsClientOutput(data, eventType) + disposition := classifyOpenAIStreamEvent(data, eventType) + startsClientOutput := forceFlushFailedEvent || disposition.commitPending + guardMayCommit := forceFlushFailedEvent || disposition.semanticOutput || disposition.timerSatisfied if guardFirstOutput { - eventStartsClientOutput = eventStartsClientOutput || startsClientOutput + eventHasSemanticOutput = eventHasSemanticOutput || disposition.semanticOutput + eventSatisfiesFirstOutput = eventSatisfiesFirstOutput || forceFlushFailedEvent || disposition.timerSatisfied } // 鍐欏叆瀹㈡埛绔紙瀹㈡埛绔柇寮€鍚庣户缁?drain 涓婃父锛? if !clientDisconnected { - shouldFlush := queueDrained && (clientOutputStarted || startsClientOutput) - if firstTokenMs == nil && startsClientOutput { + commitNow := startsClientOutput + if guardFirstOutput { + commitNow = guardMayCommit + } + shouldFlush := queueDrained && (clientOutputStarted || commitNow) + if firstTokenMs == nil && commitNow { // 淇濊瘉棣栦釜 token 浜嬩欢灏藉揩鍑虹珯锛岄伩鍏嶅奖鍝?TTFT銆? shouldFlush = true } @@ -5530,7 +6228,7 @@ func (s *OpenAIGatewayService) handleStreamingResponseWithReasoning(ctx context. } // Record first token time - if !guardFirstOutput && firstTokenMs == nil && startsClientOutput { + if !guardFirstOutput && firstTokenMs == nil && disposition.semanticOutput { ms := int(time.Since(startTime).Milliseconds()) firstTokenMs = &ms stopFirstOutputTimer() @@ -5671,7 +6369,17 @@ func (s *OpenAIGatewayService) handleStreamingResponseWithReasoning(ctx context. if clientDisconnected { return resultWithUsage(), fmt.Errorf("stream usage incomplete after timeout") } + if sawTerminalEvent && !sawFailedEvent { + _ = resp.Body.Close() + return resultWithUsage(), nil + } logger.LegacyPrintf("service.openai_gateway", "Stream data interval timeout: account=%d model=%s interval=%s", account.ID, originalModel, streamInterval) + if guardFirstOutput && !clientOutputStarted { + _ = resp.Body.Close() + failoverErr := s.newOpenAIStreamFailoverError(c, account, false, upstreamRequestID, nil, "OpenAI stream produced no complete output event before the data interval timeout") + failoverErr.SafeToFailoverAfterWrite = neutralKeepaliveWritten + return resultWithUsage(), failoverErr + } // 澶勭悊娴佽秴鏃讹紝鍙兘鏍囪璐︽埛涓轰复鏃朵笉鍙皟搴︽垨閿欒鐘舵€? if s.rateLimitService != nil { s.rateLimitService.HandleStreamTimeout(ctx, account, originalModel) @@ -5692,11 +6400,11 @@ func (s *OpenAIGatewayService) handleStreamingResponseWithReasoning(ctx context. ctx, c, account, - startTime, + firstOutputStart, originalModel, reasoningEffort, firstOutputTimeout, - "semantic_output", + openAIFirstOutputPhaseSemanticOutput, resp.Header, ) @@ -5720,6 +6428,7 @@ func (s *OpenAIGatewayService) handleStreamingResponseWithReasoning(ctx context. continue } flusher.Flush() + neutralKeepaliveWritten = true lastDownstreamWriteAt = time.Now() continue } @@ -6774,6 +7483,7 @@ func (s *OpenAIGatewayService) RecordUsage(ctx context.Context, input *OpenAIRec OutputTokens: result.Usage.OutputTokens, CacheCreationTokens: result.Usage.CacheCreationInputTokens, CacheReadTokens: result.Usage.CacheReadInputTokens, + ImageInputTokens: tokens.ImageInputTokens, ImageOutputTokens: result.Usage.ImageOutputTokens, ImageCount: result.ImageCount, ImageSize: optionalTrimmedStringPtr(result.ImageSize), @@ -6786,6 +7496,7 @@ func (s *OpenAIGatewayService) RecordUsage(ctx context.Context, input *OpenAIRec } if cost != nil { usageLog.InputCost = cost.InputCost + usageLog.ImageInputCost = cost.ImageInputCost usageLog.OutputCost = cost.OutputCost usageLog.ImageOutputCost = cost.ImageOutputCost usageLog.CacheCreationCost = cost.CacheCreationCost diff --git a/backend/internal/service/openai_gateway_service_test.go b/backend/internal/service/openai_gateway_service_test.go index 65cf1cf8a..a73cd82fd 100644 --- a/backend/internal/service/openai_gateway_service_test.go +++ b/backend/internal/service/openai_gateway_service_test.go @@ -714,7 +714,7 @@ func TestOpenAISelectAccountWithScheduler_AccountShareModeUnsupportedModelDoesNo } ctx := WithAccountShareModeRequest(context.Background(), consumerUserID, apiKeyID) - selection, decision, err := svc.selectAccountWithScheduler(ctx, &modeGroupID, "", "", "gpt-unsupported", nil, OpenAIUpstreamTransportHTTPSSE, "", false) + selection, decision, err := svc.selectAccountWithScheduler(ctx, &modeGroupID, "", "", "gpt-unsupported", nil, OpenAIUpstreamTransportHTTPSSE, "", "", false) require.Nil(t, selection) require.Equal(t, openAIAccountScheduleLayerAccountShareMode, decision.Layer) @@ -723,6 +723,90 @@ func TestOpenAISelectAccountWithScheduler_AccountShareModeUnsupportedModelDoesNo require.NotNil(t, shareRepo.membership) } +func TestOpenAIGatewayServiceRevalidateSelectedAccountForDispatchUsesLatestPublicState(t *testing.T) { + groupID := int64(7123) + ownerUserID := int64(91) + consumerUserID := int64(92) + account := Account{ + ID: 551, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, + Concurrency: 3, + OwnerUserID: &ownerUserID, + ShareMode: AccountShareModePublic, + ShareStatus: AccountShareStatusApproved, + GroupIDs: []int64{groupID}, + AccountLevel: AccountLevelTeam, + } + repo := &stubOpenAIAccountRepo{accounts: []Account{account}} + svc := &OpenAIGatewayService{accountRepo: repo} + ctx := context.WithValue(context.Background(), ctxkey.AuthenticatedUserID, consumerUserID) + + requirements := OpenAIAccountDispatchRequirements{RequestedModel: "gpt-5", RequiredTransport: OpenAIUpstreamTransportAny} + latest, err := svc.RevalidateSelectedOpenAIAccountForDispatch(ctx, &groupID, &account, requirements) + require.NoError(t, err) + require.NotNil(t, latest) + require.Equal(t, account.ID, latest.ID) + + repo.accounts[0].ShareStatus = AccountShareStatusPending + latest, err = svc.RevalidateSelectedOpenAIAccountForDispatch(ctx, &groupID, &account, requirements) + require.ErrorIs(t, err, ErrNoAvailableAccounts) + require.Nil(t, latest) + + repo.accounts[0].ShareStatus = AccountShareStatusApproved + repo.accounts[0].GroupIDs = nil + latest, err = svc.RevalidateSelectedOpenAIAccountForDispatch(ctx, &groupID, &account, requirements) + require.ErrorIs(t, err, ErrNoAvailableAccounts) + require.Nil(t, latest) +} + +func TestOpenAIGatewayServiceRevalidateSelectedAccountForDispatchUsesModeMembership(t *testing.T) { + modeGroupID := int64(8123) + ownerUserID := int64(101) + consumerUserID := int64(202) + apiKeyID := int64(303) + account := Account{ + ID: 661, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, + Concurrency: 3, + OwnerUserID: &ownerUserID, + ShareMode: AccountShareModePrivate, + ShareStatus: AccountShareStatusApproved, + GroupIDs: []int64{modeGroupID}, + AccountLevel: AccountLevelTeam, + } + shareRepo := &accountShareModeRepoStub{ + membership: &AccountShareMembership{ID: 1, AccountID: account.ID, ConsumerUserID: consumerUserID, APIKeyID: apiKeyID}, + listing: &AccountShareListing{ID: 2, AccountID: account.ID, OwnerUserID: ownerUserID, Status: AccountShareListingStatusActive}, + } + svc := &OpenAIGatewayService{ + accountRepo: stubOpenAIAccountRepo{accounts: []Account{account}}, + accountShareModeService: &AccountShareModeService{repo: shareRepo}, + } + ctx := WithAccountShareModeRequest(context.Background(), consumerUserID, apiKeyID) + + requirements := OpenAIAccountDispatchRequirements{RequestedModel: "gpt-5", RequiredTransport: OpenAIUpstreamTransportAny} + latest, err := svc.RevalidateSelectedOpenAIAccountForDispatch(ctx, &modeGroupID, &account, requirements) + require.NoError(t, err) + require.NotNil(t, latest) + require.Equal(t, account.ID, latest.ID) + + wrongRepo := &accountShareModeRepoStub{ + membership: &AccountShareMembership{ID: 3, AccountID: account.ID + 1, ConsumerUserID: consumerUserID, APIKeyID: apiKeyID}, + listing: &AccountShareListing{ID: 4, AccountID: account.ID + 1, OwnerUserID: ownerUserID, Status: AccountShareListingStatusActive}, + } + svc.accountShareModeService = &AccountShareModeService{repo: wrongRepo} + ctx = WithAccountShareModeRequest(context.Background(), consumerUserID, apiKeyID) + latest, err = svc.RevalidateSelectedOpenAIAccountForDispatch(ctx, &modeGroupID, &account, requirements) + require.ErrorIs(t, err, ErrAccountShareModeGroupUnbound) + require.Nil(t, latest) +} + func TestOpenAISelectAccountWithLoadAwareness_FiltersUnschedulableWhenNoConcurrencyService(t *testing.T) { now := time.Now() resetAt := now.Add(10 * time.Minute) diff --git a/backend/internal/service/openai_grok_selection.go b/backend/internal/service/openai_grok_selection.go index 4991b96dc..e4c5c8592 100644 --- a/backend/internal/service/openai_grok_selection.go +++ b/backend/internal/service/openai_grok_selection.go @@ -35,6 +35,7 @@ func (s *OpenAIGatewayService) SelectAccountWithSchedulerForGrok( sessionHash string, requestedModel string, excludedIDs map[int64]struct{}, + requiredEndpointCapability OpenAIEndpointCapability, ) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error) { ctx = withGrokPlatform(ctx) selection, decision, err := s.selectAccountWithScheduler( @@ -46,12 +47,19 @@ func (s *OpenAIGatewayService) SelectAccountWithSchedulerForGrok( excludedIDs, OpenAIUpstreamTransportHTTPSSE, "", + requiredEndpointCapability, false, ) if err != nil || selection == nil || selection.Account == nil { return selection, decision, err } if selection.Account.Platform == PlatformGrok { + selection.OpenAIDispatchRequirements = &OpenAIAccountDispatchRequirements{ + RequestedModel: requestedModel, + RequiredTransport: OpenAIUpstreamTransportHTTPSSE, + RequiredEndpointCapability: requiredEndpointCapability, + RequiredPlatform: PlatformGrok, + } return selection, decision, nil } if selection.ReleaseFunc != nil { diff --git a/backend/internal/service/openai_images.go b/backend/internal/service/openai_images.go index b58cadbbe..c370344b4 100644 --- a/backend/internal/service/openai_images.go +++ b/backend/internal/service/openai_images.go @@ -645,12 +645,14 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesAPIKey( Kind: "failover", Message: upstreamMsg, }) - s.handleFailoverSideEffects(upstreamCtx, resp, account) - return nil, &UpstreamFailoverError{ - StatusCode: resp.StatusCode, - ResponseBody: respBody, - RetryableOnSameAccount: shouldRetryOpenAIOnSamePoolAccount(account, resp.StatusCode, upstreamMsg, respBody), - } + s.handleFailoverSideEffectsForModel(upstreamCtx, resp, account, requestModel) + return nil, newOpenAIUpstreamFailoverError( + resp.StatusCode, + resp.Header, + respBody, + upstreamMsg, + shouldRetryOpenAIOnSamePoolAccount(account, resp.StatusCode, upstreamMsg, respBody), + ) } if resp.StatusCode == http.StatusBadRequest { return nil, newOpenAIImagesStreamFailoverError(resp, resp.StatusCode, upstreamMsg, false) diff --git a/backend/internal/service/openai_images_responses.go b/backend/internal/service/openai_images_responses.go index 69be6057a..14d84efbd 100644 --- a/backend/internal/service/openai_images_responses.go +++ b/backend/internal/service/openai_images_responses.go @@ -1057,12 +1057,14 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesOAuth( Kind: "failover", Message: upstreamMsg, }) - s.handleFailoverSideEffects(upstreamCtx, resp, account) - return nil, &UpstreamFailoverError{ - StatusCode: resp.StatusCode, - ResponseBody: respBody, - RetryableOnSameAccount: shouldRetryOpenAIOnSamePoolAccount(account, resp.StatusCode, upstreamMsg, respBody), - } + s.handleFailoverSideEffectsForModel(upstreamCtx, resp, account, requestModel) + return nil, newOpenAIUpstreamFailoverError( + resp.StatusCode, + resp.Header, + respBody, + upstreamMsg, + shouldRetryOpenAIOnSamePoolAccount(account, resp.StatusCode, upstreamMsg, respBody), + ) } if resp.StatusCode == http.StatusBadRequest { return nil, newOpenAIImagesStreamFailoverError(resp, resp.StatusCode, upstreamMsg, false) diff --git a/backend/internal/service/openai_request_body_limit_failover.go b/backend/internal/service/openai_request_body_limit_failover.go new file mode 100644 index 000000000..e90300f01 --- /dev/null +++ b/backend/internal/service/openai_request_body_limit_failover.go @@ -0,0 +1,43 @@ +package service + +import "net/http" + +// OpenAIRequestBodyTooLargeClientMessage is deliberately provider-neutral: +// upstream bodies can contain proxy, host, credential, or account details. +const OpenAIRequestBodyTooLargeClientMessage = "Request payload is too large" + +const openAIRequestBodyTooLargeReason = GatewayFailureReason("openai_request_body_too_large") + +func isOpenAIRequestBodyTooLargeError(statusCode int, upstreamMsg string, upstreamBody []byte) bool { + return statusCode == http.StatusRequestEntityTooLarge && !isOpenAIContextWindowError(upstreamMsg, upstreamBody) +} + +func newOpenAIUpstreamFailoverError( + statusCode int, + responseHeaders http.Header, + responseBody []byte, + upstreamMsg string, + retryableOnSameAccount bool, +) *UpstreamFailoverError { + failoverErr := &UpstreamFailoverError{ + StatusCode: statusCode, + ResponseBody: responseBody, + ResponseHeaders: responseHeaders.Clone(), + RetryableOnSameAccount: retryableOnSameAccount, + } + if isOpenAIRequestBodyTooLargeError(statusCode, upstreamMsg, responseBody) { + failoverErr.RetryableOnSameAccount = false + failoverErr.Scope = GatewayFailureScopeAccount + failoverErr.Reason = openAIRequestBodyTooLargeReason + failoverErr.NextAccountAction = NextAccountRetry + failoverErr.ClientStatusCode = http.StatusRequestEntityTooLarge + failoverErr.ClientMessage = OpenAIRequestBodyTooLargeClientMessage + } + return failoverErr +} + +// IsOpenAIRequestBodyTooLarge reports an account-specific serialized-body +// limit. The same request may still fit another compatible upstream account. +func (e *UpstreamFailoverError) IsOpenAIRequestBodyTooLarge() bool { + return e != nil && e.Reason == openAIRequestBodyTooLargeReason +} diff --git a/backend/internal/service/openai_request_body_limit_failover_test.go b/backend/internal/service/openai_request_body_limit_failover_test.go new file mode 100644 index 000000000..21c9db8e1 --- /dev/null +++ b/backend/internal/service/openai_request_body_limit_failover_test.go @@ -0,0 +1,43 @@ +package service + +import ( + "net/http" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestNewOpenAIUpstreamFailoverErrorClassifiesAccountBodyLimit(t *testing.T) { + headers := http.Header{"X-Request-Id": []string{"req-body-limit"}} + body := []byte(`{"error":{"type":"invalid_request_error","message":"request body exceeds the maximum allowed size"}}`) + + failoverErr := newOpenAIUpstreamFailoverError( + http.StatusRequestEntityTooLarge, + headers, + body, + "request body exceeds the maximum allowed size", + true, + ) + + require.True(t, failoverErr.IsOpenAIRequestBodyTooLarge()) + require.Equal(t, GatewayFailureScopeAccount, failoverErr.Scope) + require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction) + require.False(t, failoverErr.RetryableOnSameAccount) + require.Equal(t, http.StatusRequestEntityTooLarge, failoverErr.ClientStatusCode) + require.Equal(t, OpenAIRequestBodyTooLargeClientMessage, failoverErr.ClientMessage) + require.Equal(t, "req-body-limit", failoverErr.ResponseHeaders.Get("X-Request-Id")) +} + +func TestOpenAIRequestBodyLimitDoesNotReclassifyContextWindowError(t *testing.T) { + body := []byte(`{"error":{"code":"context_length_exceeded","message":"maximum context length exceeded"}}`) + failoverErr := newOpenAIUpstreamFailoverError( + http.StatusRequestEntityTooLarge, + nil, + body, + "maximum context length exceeded", + false, + ) + + require.False(t, failoverErr.IsOpenAIRequestBodyTooLarge()) + require.False(t, shouldFailoverOpenAIPassthroughResponse(http.StatusRequestEntityTooLarge, "maximum context length exceeded", body)) +} diff --git a/backend/internal/service/openai_responses_lite_tools.go b/backend/internal/service/openai_responses_lite_tools.go index a15bc498b..c5c16ecaa 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 173dc370b..18ad5725e 100644 --- a/backend/internal/service/openai_responses_lite_tools_test.go +++ b/backend/internal/service/openai_responses_lite_tools_test.go @@ -79,6 +79,43 @@ func TestNormalizeOpenAIResponsesLiteToolsRejectsUnsupportedHostedTool(t *testin require.False(t, changed) } +func TestNormalizeOpenAIResponsesLiteToolsEnsuresAllTurnsReasoningContext(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"]) + }) + } +} + +func TestNormalizeOpenAIResponsesLiteToolsRejectsNonObjectReasoning(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 TestCodexImageFunctionToolPreventsNativeImageToolInjection(t *testing.T) { req := map[string]any{ "model": "gpt-5.4", 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 000000000..d8d2bb02b --- /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 000000000..d922d1d46 --- /dev/null +++ b/backend/internal/service/openai_responses_rejected_field_retry_test.go @@ -0,0 +1,160 @@ +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/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint" + "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: "parameter is invalid but not unsupported", + 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 message", + body: []byte(`{"input":[{"type":"function_call","namespace":"keep","arguments":"{}"}]}`), + responseBody: []byte(`{"error":{"code":"unknown_parameter","message":"Unknown parameter: input[0].namespace","param":"tools"}}`), + }, + } + + 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 TestOpenAIGatewayServiceRetriesExplicitlyRejectedResponsesFields(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 := &openAIRejectedFieldUpstream{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 *openAIRejectedFieldUpstream) *OpenAIGatewayService { + return &OpenAIGatewayService{ + cfg: &config.Config{Security: config.SecurityConfig{ + URLAllowlist: config.URLAllowlistConfig{Enabled: false}, + }}, + httpUpstream: upstream, + } +} + +type openAIRejectedFieldUpstream struct { + responses []*http.Response + bodies [][]byte +} + +func (u *openAIRejectedFieldUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) { + body, err := io.ReadAll(req.Body) + if err != nil { + return nil, err + } + u.bodies = append(u.bodies, body) + if len(u.responses) == 0 { + return nil, fmt.Errorf("unexpected upstream request %d", len(u.bodies)) + } + response := u.responses[0] + u.responses = u.responses[1:] + return response, nil +} + +func (u *openAIRejectedFieldUpstream) DoWithTLS(req *http.Request, proxyURL string, accountID int64, accountConcurrency int, _ *tlsfingerprint.Profile) (*http.Response, error) { + return u.Do(req, proxyURL, accountID, accountConcurrency) +} + +func newOpenAIRejectedFieldTestContext(body []byte) *gin.Context { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + c.Request.Header.Set("User-Agent", "curl/8.0") + return c +} + +func newOpenAIRejectedFieldTestAccount() *Account { + return &Account{ + ID: 5107, + Name: "responses-compatible", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-test", + "base_url": "https://compat.example", + }, + Extra: map[string]any{ + openai_compat.ExtraKeyResponsesMode: string(openai_compat.ResponsesSupportModeAuto), + openai_compat.ExtraKeyResponsesSupported: true, + }, + Status: StatusActive, + Schedulable: true, + } +} + +func newOpenAIRejectedFieldTestResponse(status int, body string) *http.Response { + return &http.Response{ + StatusCode: status, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(body)), + } +} diff --git a/backend/internal/service/openai_ws_client_read.go b/backend/internal/service/openai_ws_client_read.go new file mode 100644 index 000000000..d76a74016 --- /dev/null +++ b/backend/internal/service/openai_ws_client_read.go @@ -0,0 +1,71 @@ +package service + +import ( + "context" + "errors" + "time" + + coderws "github.com/coder/websocket" +) + +type openAIWSClientReadResult struct { + messageType coderws.MessageType + payload []byte + err error +} + +// ReadOpenAIWSClientMessage keeps a single read goroutine alive while a +// control event sends its close frame, then closes the transport and joins the +// reader. This prevents canceled ingress reads from lingering after teardown. +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 timer *time.Timer + var timeoutCh <-chan time.Time + 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_forwarder.go b/backend/internal/service/openai_ws_forwarder.go index eaf29f882..72dd55f8c 100644 --- a/backend/internal/service/openai_ws_forwarder.go +++ b/backend/internal/service/openai_ws_forwarder.go @@ -217,6 +217,13 @@ func (e *OpenAIWSClientCloseError) Reason() string { return strings.TrimSpace(e.reason) } +func openAIWSImageGenerationPermissionError(imageGenerationAllowed bool, requestedModel string, payload []byte) error { + if imageGenerationAllowed || !IsExplicitImageGenerationIntent(openAIResponsesEndpoint, requestedModel, payload) { + return nil + } + return NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, ImageGenerationPermissionMessage(), nil) +} + // OpenAIWSIngressHooks 定义入站 WS 每个 turn 的生命周期回调。 type OpenAIWSIngressHooks struct { BeforeTurn func(turn int) error @@ -2105,7 +2112,7 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( return nil, &agentIdentityTaskRecoveredError{} } if errors.As(err, &dialErr) && dialErr != nil && dialErr.StatusCode == http.StatusTooManyRequests { - s.persistOpenAIWSRateLimitSignal(ctx, account, dialErr.ResponseHeaders, nil, "rate_limit_exceeded", "rate_limit_error", strings.TrimSpace(err.Error())) + s.persistOpenAIWSRateLimitSignal(ctx, account, originalModel, dialErr.ResponseHeaders, nil, "rate_limit_exceeded", "rate_limit_error", strings.TrimSpace(err.Error())) } return nil, wrapOpenAIWSFallback(classifyOpenAIWSAcquireError(err), err) } @@ -2183,6 +2190,7 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( lease, decision, payload, + originalModel, previousResponseID, needsToolContinuation, account, @@ -2336,6 +2344,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) @@ -2364,6 +2391,9 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( setOpsUpstreamError(c, 0, sanitizeUpstreamErrorMessage(readErr.Error()), "") return nil, fmt.Errorf("openai ws read event: %w", readErr) } + if normalized, changed := normalizeCompletedImageGenerationStatus(message); changed { + message = normalized + } eventType, eventResponseID, responseField := parseOpenAIWSEventEnvelope(message) if eventType == "" { @@ -2422,7 +2452,7 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( if eventType == "error" { errCodeRaw, errTypeRaw, errMsgRaw := parseOpenAIWSErrorEventFields(message) - s.persistOpenAIWSRateLimitSignal(ctx, account, lease.HandshakeHeaders(), message, errCodeRaw, errTypeRaw, errMsgRaw) + s.persistOpenAIWSRateLimitSignal(ctx, account, originalModel, lease.HandshakeHeaders(), message, errCodeRaw, errTypeRaw, errMsgRaw) errMsg := strings.TrimSpace(errMsgRaw) if errMsg == "" { errMsg = "Upstream websocket error" @@ -2493,7 +2523,7 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( if failedMsg == "" { failedMsg = "OpenAI model capacity temporarily unavailable" } - if s.handleOpenAIModelCapacitySignal(ctx, account, http.StatusServiceUnavailable, lease.HandshakeHeaders(), message, failedMsg) { + if s.handleOpenAIWSModelCapacitySignal(ctx, account, originalModel, http.StatusServiceUnavailable, lease.HandshakeHeaders(), message, failedMsg) { lease.MarkBroken() if !wroteDownstream { return nil, wrapOpenAIWSFallback("upstream_capacity", errors.New(failedMsg)) @@ -2892,9 +2922,11 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( logOpenAIWSModeInfo("ingress_ws_codex_spark_image_tool_stripped account_id=%d", account.ID) } } - imageIntent := IsImageGenerationIntent(openAIResponsesEndpoint, originalModel, normalized) - if imageIntent && !imageGenerationAllowed { - return openAIWSClientPayload{}, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, ImageGenerationPermissionMessage(), nil) + // Passive Codex image_gen declarations are capability catalogs, not a + // request to generate an image. Only explicit intent may hit the group + // permission gate. + if permissionErr := openAIWSImageGenerationPermissionError(imageGenerationAllowed, originalModel, normalized); permissionErr != nil { + return openAIWSClientPayload{}, permissionErr } // Apply OpenAI Fast Policy on the response.create frame using the same @@ -3129,7 +3161,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( ) var dialErr *openAIWSDialError if errors.As(acquireErr, &dialErr) && dialErr != nil && dialErr.StatusCode == http.StatusTooManyRequests { - s.persistOpenAIWSRateLimitSignal(ctx, account, dialErr.ResponseHeaders, nil, "rate_limit_exceeded", "rate_limit_error", strings.TrimSpace(acquireErr.Error())) + s.persistOpenAIWSRateLimitSignal(ctx, account, "", dialErr.ResponseHeaders, nil, "rate_limit_exceeded", "rate_limit_error", strings.TrimSpace(acquireErr.Error())) } if errors.Is(acquireErr, errOpenAIWSPreferredConnUnavailable) { return nil, NewOpenAIWSClientCloseError( @@ -3264,6 +3296,17 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( wroteDownstream, ) } + if !json.Valid(upstreamMessage) { + lease.MarkBroken() + return nil, wrapOpenAIWSIngressTurnError( + "invalid_event_json", + errors.New("upstream websocket returned invalid JSON"), + wroteDownstream, + ) + } + if normalized, changed := normalizeCompletedImageGenerationStatus(upstreamMessage); changed { + upstreamMessage = normalized + } eventType, eventResponseID, _ := parseOpenAIWSEventEnvelope(upstreamMessage) if responseID == "" && eventResponseID != "" { @@ -3279,7 +3322,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( } if eventType == "error" { errCodeRaw, errTypeRaw, errMsgRaw := parseOpenAIWSErrorEventFields(upstreamMessage) - s.persistOpenAIWSRateLimitSignal(ctx, account, lease.HandshakeHeaders(), upstreamMessage, errCodeRaw, errTypeRaw, errMsgRaw) + s.persistOpenAIWSRateLimitSignal(ctx, account, originalModel, lease.HandshakeHeaders(), upstreamMessage, errCodeRaw, errTypeRaw, errMsgRaw) fallbackReason, _ := classifyOpenAIWSErrorEventFromRaw(errCodeRaw, errTypeRaw, errMsgRaw) errCode, errType, errMessage := summarizeOpenAIWSErrorEventFieldsFromRaw(errCodeRaw, errTypeRaw, errMsgRaw) recoverablePrevNotFound := fallbackReason == openAIWSIngressStagePreviousResponseNotFound && @@ -3343,7 +3386,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( if failedMsg == "" { failedMsg = "OpenAI model capacity temporarily unavailable" } - if s.handleOpenAIModelCapacitySignal(ctx, account, http.StatusServiceUnavailable, lease.HandshakeHeaders(), upstreamMessage, failedMsg) { + if s.handleOpenAIWSModelCapacitySignal(ctx, account, originalModel, http.StatusServiceUnavailable, lease.HandshakeHeaders(), upstreamMessage, failedMsg) { lease.MarkBroken() if !wroteDownstream { return nil, NewOpenAIWSClientCloseError(coderws.StatusTryAgainLater, failedMsg, errors.New(failedMsg)) @@ -4076,6 +4119,7 @@ func (s *OpenAIGatewayService) performOpenAIWSGeneratePrewarm( lease *openAIWSConnLease, decision OpenAIWSProtocolDecision, payload map[string]any, + requestedModel string, previousResponseID string, needsToolContinuation bool, account *Account, @@ -4182,7 +4226,7 @@ func (s *OpenAIGatewayService) performOpenAIWSGeneratePrewarm( if eventType == "error" { errCodeRaw, errTypeRaw, errMsgRaw := parseOpenAIWSErrorEventFields(message) - s.persistOpenAIWSRateLimitSignal(ctx, account, lease.HandshakeHeaders(), message, errCodeRaw, errTypeRaw, errMsgRaw) + s.persistOpenAIWSRateLimitSignal(ctx, account, requestedModel, lease.HandshakeHeaders(), message, errCodeRaw, errTypeRaw, errMsgRaw) errMsg := strings.TrimSpace(errMsgRaw) if errMsg == "" { errMsg = "OpenAI websocket prewarm error" @@ -4212,7 +4256,7 @@ func (s *OpenAIGatewayService) performOpenAIWSGeneratePrewarm( if failedMsg == "" { failedMsg = "OpenAI websocket prewarm failed" } - if s.handleOpenAIModelCapacitySignal(ctx, account, http.StatusServiceUnavailable, lease.HandshakeHeaders(), message, failedMsg) { + if s.handleOpenAIWSModelCapacitySignal(ctx, account, requestedModel, http.StatusServiceUnavailable, lease.HandshakeHeaders(), message, failedMsg) { lease.MarkBroken() logOpenAIWSModeInfo( "prewarm_capacity_event account_id=%d conn_id=%s idx=%d message=%s", @@ -4492,15 +4536,40 @@ func isOpenAIWSRateLimitError(codeRaw, errTypeRaw, msgRaw string) bool { return false } -func (s *OpenAIGatewayService) persistOpenAIWSRateLimitSignal(ctx context.Context, account *Account, headers http.Header, responseBody []byte, codeRaw, errTypeRaw, msgRaw string) { - if s == nil || s.rateLimitService == nil || account == nil || account.Platform != PlatformOpenAI { +func (s *OpenAIGatewayService) handleOpenAIWSModelCapacitySignal(ctx context.Context, account *Account, requestedModel string, statusCode int, headers http.Header, responseBody []byte, message string) bool { + if s == nil || account == nil || account.Platform != PlatformOpenAI { + return false + } + if statusCode <= 0 { + statusCode = http.StatusServiceUnavailable + } + if !isOpenAITransientCapacityError(statusCode, message, responseBody) { + return false + } + cooldownBody := responseBody + if len(cooldownBody) == 0 { + cooldownBody = []byte(message) + } + s.handleOpenAIAccountUpstreamErrorForModel(ctx, account, requestedModel, statusCode, headers, cooldownBody) + return true +} + +func (s *OpenAIGatewayService) persistOpenAIWSRateLimitSignal(ctx context.Context, account *Account, requestedModel string, headers http.Header, responseBody []byte, codeRaw, errTypeRaw, msgRaw string) { + if s == nil || account == nil || account.Platform != PlatformOpenAI { return } - if isOpenAITransientCapacityError(http.StatusServiceUnavailable, strings.TrimSpace(msgRaw+" "+codeRaw+" "+errTypeRaw), responseBody) { - s.rateLimitService.HandleUpstreamError(ctx, account, http.StatusServiceUnavailable, headers, responseBody) + if s.handleOpenAIWSModelCapacitySignal( + ctx, + account, + requestedModel, + http.StatusServiceUnavailable, + headers, + responseBody, + strings.TrimSpace(msgRaw+" "+codeRaw+" "+errTypeRaw), + ) { return } - if !isOpenAIWSRateLimitError(codeRaw, errTypeRaw, msgRaw) { + if s.rateLimitService == nil || !isOpenAIWSRateLimitError(codeRaw, errTypeRaw, msgRaw) { return } s.rateLimitService.HandleUpstreamError(ctx, account, http.StatusTooManyRequests, headers, responseBody) 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 9bf77aa79..0bf7ca348 100644 --- a/backend/internal/service/openai_ws_forwarder_ingress_session_test.go +++ b/backend/internal/service/openai_ws_forwarder_ingress_session_test.go @@ -38,6 +38,7 @@ func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_KeepLeaseAcrossT captureConn := &openAIWSCaptureConn{ events: [][]byte{ + []byte(`{"type":"response.output_item.done","item":{"id":"ig_ingress_1","type":"image_generation_call","status":"generating","result":"iVBORw0KGgoAAAANSUhEUg/+=="}}`), []byte(`{"type":"response.completed","response":{"id":"resp_ingress_turn_1","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`), []byte(`{"type":"response.completed","response":{"id":"resp_ingress_turn_2","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`), }, @@ -98,6 +99,7 @@ func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_KeepLeaseAcrossT req.Header = req.Header.Clone() req.Header.Set("User-Agent", "unit-test-agent/1.0") ginCtx.Request = req + ginCtx.Set("api_key", &APIKey{Group: &Group{AllowImageGeneration: false}}) readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second) msgType, firstMessage, readErr := conn.Read(readCtx) @@ -137,7 +139,11 @@ func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_KeepLeaseAcrossT return message } - writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false}`) + writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"tools":[{"type":"namespace","name":"image_gen","tools":[{"type":"function","name":"imagegen"}]},{"type":"function","name":"image_gen.imagegen"}],"tool_choice":"auto"}`) + firstTurnImageEvent := readMessage() + require.Equal(t, "response.output_item.done", gjson.GetBytes(firstTurnImageEvent, "type").String()) + require.Equal(t, "completed", gjson.GetBytes(firstTurnImageEvent, "item.status").String()) + require.Equal(t, "iVBORw0KGgoAAAANSUhEUg/+==", gjson.GetBytes(firstTurnImageEvent, "item.result").String()) firstTurnEvent := readMessage() require.Equal(t, "response.completed", gjson.GetBytes(firstTurnEvent, "type").String()) require.Equal(t, "resp_ingress_turn_1", gjson.GetBytes(firstTurnEvent, "response.id").String()) @@ -164,6 +170,193 @@ func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_KeepLeaseAcrossT require.Len(t, captureConn.writes, 2, "应向同一上游连接发送两轮 response.create") } +func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_RejectsMalformedCtxPoolUpstreamJSON(t *testing.T) { + gin.SetMode(gin.TestMode) + + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + cfg.Gateway.OpenAIWS.Enabled = true + cfg.Gateway.OpenAIWS.OAuthEnabled = true + cfg.Gateway.OpenAIWS.APIKeyEnabled = true + cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true + cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1 + cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0 + cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1 + cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8 + cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3 + + upstreamConn := &openAIWSCaptureConn{events: [][]byte{[]byte(`{"type":"response.completed"`)}} + dialer := &openAIWSCaptureDialer{conn: upstreamConn} + pool := newOpenAIWSConnPool(cfg) + pool.setClientDialerForTest(dialer) + svc := &OpenAIGatewayService{ + cfg: cfg, + httpUpstream: &httpUpstreamRecorder{}, + cache: &stubGatewayCache{}, + openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), + toolCorrector: NewCodexToolCorrector(), + openaiWSPool: pool, + } + account := &Account{ + ID: 214, + Name: "openai-ingress-invalid-json", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Credentials: map[string]any{"api_key": "sk-test"}, + Extra: map[string]any{"responses_websockets_v2_enabled": true}, + } + + serverErrCh := make(chan error, 1) + afterTurnErrCh := make(chan error, 1) + wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := coderws.Accept(w, r, nil) + if err != nil { + serverErrCh <- err + return + } + defer func() { _ = conn.CloseNow() }() + ginCtx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ginCtx.Request = r.Clone(r.Context()) + readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second) + _, firstMessage, readErr := conn.Read(readCtx) + cancel() + if readErr != nil { + serverErrCh <- readErr + return + } + hooks := &OpenAIWSIngressHooks{AfterTurn: func(_ int, _ *OpenAIForwardResult, turnErr error) { + if turnErr != nil { + afterTurnErrCh <- turnErr + } + }} + serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, hooks) + })) + defer wsServer.Close() + + dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second) + clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil) + cancelDial() + require.NoError(t, err) + defer func() { _ = clientConn.CloseNow() }() + writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second) + require.NoError(t, clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1"}`))) + cancelWrite() + readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second) + _, _, readErr := clientConn.Read(readCtx) + cancelRead() + require.Error(t, readErr, "malformed upstream event must close the client path without forwarding") + + serverErr := <-serverErrCh + require.ErrorContains(t, serverErr, "invalid JSON") + afterTurnErr := <-afterTurnErrCh + require.ErrorContains(t, afterTurnErr, "invalid JSON") + upstreamConn.mu.Lock() + closed := upstreamConn.closed + upstreamConn.mu.Unlock() + require.True(t, closed, "malformed upstream connection must be marked broken and closed") + require.Equal(t, 1, dialer.DialCount(), "malformed event must fail immediately rather than retry") +} + +func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_RejectsExplicitImageIntent(t *testing.T) { + runOpenAIWSExplicitImagePermissionRejection(t, false) +} + +func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_PassthroughRejectsExplicitImageIntent(t *testing.T) { + runOpenAIWSExplicitImagePermissionRejection(t, true) +} + +func runOpenAIWSExplicitImagePermissionRejection(t *testing.T, passthrough bool) { + t.Helper() + gin.SetMode(gin.TestMode) + + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + cfg.Gateway.OpenAIWS.Enabled = true + cfg.Gateway.OpenAIWS.APIKeyEnabled = true + cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true + cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = passthrough + + accountExtra := map[string]any{"responses_websockets_v2_enabled": true} + if passthrough { + accountExtra["openai_apikey_responses_websockets_v2_mode"] = OpenAIWSIngressModePassthrough + } + + svc := &OpenAIGatewayService{ + cfg: cfg, + httpUpstream: &httpUpstreamRecorder{}, + cache: &stubGatewayCache{}, + openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), + toolCorrector: NewCodexToolCorrector(), + } + account := &Account{ + ID: 115, + Name: "openai-ingress-image-permission", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Credentials: map[string]any{"api_key": "sk-test"}, + Extra: accountExtra, + } + + serverErrCh := make(chan error, 1) + wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := coderws.Accept(w, r, nil) + if err != nil { + serverErrCh <- err + return + } + defer func() { _ = conn.CloseNow() }() + + ginCtx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ginCtx.Request = r.Clone(r.Context()) + ginCtx.Set("api_key", &APIKey{Group: &Group{AllowImageGeneration: false}}) + + readCtx, cancelRead := context.WithTimeout(r.Context(), 3*time.Second) + msgType, firstMessage, readErr := conn.Read(readCtx) + cancelRead() + if readErr != nil { + serverErrCh <- readErr + return + } + if msgType != coderws.MessageText { + serverErrCh <- errors.New("unsupported websocket client message type") + return + } + serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil) + })) + defer wsServer.Close() + + dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second) + clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil) + cancelDial() + require.NoError(t, err) + defer func() { _ = clientConn.CloseNow() }() + + writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second) + err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1","tools":[{"type":"function","name":"image_gen.imagegen"}],"tool_choice":{"type":"function","name":"image_gen.imagegen"},"input":"draw"}`)) + cancelWrite() + require.NoError(t, err) + + select { + case serverErr := <-serverErrCh: + var closeErr *OpenAIWSClientCloseError + require.ErrorAs(t, serverErr, &closeErr) + require.Equal(t, coderws.StatusPolicyViolation, closeErr.statusCode) + require.Equal(t, ImageGenerationPermissionMessage(), closeErr.reason) + case <-time.After(5 * time.Second): + t.Fatal("等待显式生图权限拒绝超时") + } +} + func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_RepairsStalePreviousResponseID(t *testing.T) { gin.SetMode(gin.TestMode) @@ -453,8 +646,8 @@ func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_PassthroughModeR upstreamConn := &openAIWSCaptureConn{ events: [][]byte{ - []byte(`{"type":"response.output_item.done","item":{"id":"ig_passthrough_turn_1","type":"image_generation_call","result":"ZmluYWw="}}`), - []byte(`{"type":"response.completed","response":{"id":"resp_passthrough_turn_1","model":"gpt-5.1","usage":{"input_tokens":2,"output_tokens":3,"output_tokens_details":{"image_tokens":2}},"output":[{"id":"ig_passthrough_turn_1","type":"image_generation_call","result":"ZmluYWw="}]}}`), + []byte(`{"type":"response.output_item.done","item":{"id":"ig_passthrough_turn_1","type":"image_generation_call","status":"generating","result":"ZmluYWw="}}`), + []byte(`{"type":"response.completed","response":{"id":"resp_passthrough_turn_1","model":"gpt-5.1","usage":{"input_tokens":2,"output_tokens":3,"output_tokens_details":{"image_tokens":2}},"output":[{"id":"ig_passthrough_turn_1","type":"image_generation_call","status":"in_progress","result":"ZmluYWw="}]}}`), }, } captureDialer := &openAIWSCaptureDialer{conn: upstreamConn} @@ -557,12 +750,14 @@ func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_PassthroughModeR cancelRead() require.NoError(t, readErr) require.Equal(t, "response.output_item.done", gjson.GetBytes(event, "type").String()) + require.Equal(t, "completed", gjson.GetBytes(event, "item.status").String()) readCtx, cancelRead = context.WithTimeout(context.Background(), 3*time.Second) _, event, readErr = clientConn.Read(readCtx) cancelRead() require.NoError(t, readErr) require.Equal(t, "response.completed", gjson.GetBytes(event, "type").String()) require.Equal(t, "resp_passthrough_turn_1", gjson.GetBytes(event, "response.id").String()) + require.Equal(t, "completed", gjson.GetBytes(event, "response.output.0.status").String()) _ = clientConn.Close(coderws.StatusNormalClosure, "done") select { diff --git a/backend/internal/service/openai_ws_forwarder_success_test.go b/backend/internal/service/openai_ws_forwarder_success_test.go index d812698c3..eee5e29ee 100644 --- a/backend/internal/service/openai_ws_forwarder_success_test.go +++ b/backend/internal/service/openai_ws_forwarder_success_test.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "errors" + "fmt" "io" "net/http" "net/http/httptest" @@ -72,6 +73,12 @@ func TestOpenAIGatewayService_Forward_WSv2_SuccessAndBindSticky(t *testing.T) { "response": map[string]any{ "id": "resp_new_1", "model": "gpt-5.1", + "output": []map[string]any{{ + "id": "ig_ws_1", + "type": "image_generation_call", + "status": "generating", + "result": "final-image", + }}, "usage": map[string]any{ "input_tokens": 12, "output_tokens": 7, @@ -171,6 +178,7 @@ func TestOpenAIGatewayService_Forward_WSv2_SuccessAndBindSticky(t *testing.T) { responseBody := rec.Body.Bytes() require.Equal(t, "resp_new_1", gjson.GetBytes(responseBody, "id").String()) + require.Equal(t, "completed", gjson.GetBytes(responseBody, "output.0.status").String()) } func requestToJSONString(payload map[string]any) string { @@ -864,7 +872,7 @@ func TestOpenAIGatewayService_Forward_WSv2_GeneratePrewarm(t *testing.T) { require.False(t, gjson.Get(secondWrite, "generate").Exists()) } -func TestOpenAIGatewayService_GeneratePrewarmResponseFailedCapacityTempUnscheds(t *testing.T) { +func TestOpenAIGatewayService_GeneratePrewarmResponseFailedCapacityUsesModelScopedCooldown(t *testing.T) { cfg := &config.Config{} cfg.Gateway.OpenAIWS.PrewarmGenerateEnabled = true cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3 @@ -885,39 +893,45 @@ func TestOpenAIGatewayService_GeneratePrewarmResponseFailedCapacityTempUnscheds( Schedulable: true, Extra: map[string]any{"pool_mode": true}, } - conn := newOpenAIWSConn("prewarm_capacity_conn", account.ID, &openAIWSCaptureConn{ - events: [][]byte{ - []byte(`{"type":"response.failed","error":{"code":"model_capacity_exhausted","message":"Selected model is at capacity. Please try a different model."}}`), - }, - }, nil) - lease := &openAIWSConnLease{ - accountID: account.ID, - conn: conn, + newLease := func(connID string) *openAIWSConnLease { + conn := newOpenAIWSConn(connID, account.ID, &openAIWSCaptureConn{ + events: [][]byte{ + []byte(`{"type":"response.failed","error":{"code":"model_capacity_exhausted","message":"Selected model is at capacity. Please try a different model."}}`), + }, + }, nil) + return &openAIWSConnLease{ + accountID: account.ID, + conn: conn, + } } payload := map[string]any{ "type": "response.create", "model": "gpt-5.1", } - start := time.Now() - err := svc.performOpenAIWSGeneratePrewarm( - context.Background(), - lease, - OpenAIWSProtocolDecision{Transport: OpenAIUpstreamTransportResponsesWebsocketV2}, - payload, - "", - false, - account, - nil, - 0, - ) + for i := range 2 { + lease := newLease(fmt.Sprintf("prewarm_capacity_conn_%d", i)) + err := svc.performOpenAIWSGeneratePrewarm( + context.Background(), + lease, + OpenAIWSProtocolDecision{Transport: OpenAIUpstreamTransportResponsesWebsocketV2}, + payload, + "gpt-5.1", + "", + false, + account, + nil, + 0, + ) + + require.Error(t, err) + require.Contains(t, err.Error(), "prewarm_upstream_capacity") + require.False(t, lease.IsPrewarmed()) + } - require.Error(t, err) - require.Contains(t, err.Error(), "prewarm_upstream_capacity") - require.False(t, lease.IsPrewarmed()) - require.Len(t, repo.tempCalls, 1) - require.WithinDuration(t, start.Add(openAIModelCapacityCooldown), repo.tempCalls[0], 5*time.Second) - require.Contains(t, repo.tempReasons[0], "openai_model_capacity") + require.Empty(t, repo.tempCalls, "API Key 瞬态容量错误不得冷却整个账号") + require.True(t, svc.isOpenAIAccountModelRuntimeBlocked(account, "gpt-5.1")) + require.False(t, svc.isOpenAIAccountModelRuntimeBlocked(account, "gpt-5.2")) } func TestOpenAIGatewayService_PrewarmReadHonorsParentContext(t *testing.T) { @@ -958,6 +972,7 @@ func TestOpenAIGatewayService_PrewarmReadHonorsParentContext(t *testing.T) { lease, OpenAIWSProtocolDecision{Transport: OpenAIUpstreamTransportResponsesWebsocketV2}, payload, + "gpt-5.1", "", false, account, diff --git a/backend/internal/service/openai_ws_ratelimit_signal_test.go b/backend/internal/service/openai_ws_ratelimit_signal_test.go index ba3e723cb..33c0adf6a 100644 --- a/backend/internal/service/openai_ws_ratelimit_signal_test.go +++ b/backend/internal/service/openai_ws_ratelimit_signal_test.go @@ -262,7 +262,7 @@ func TestOpenAIGatewayService_WSv2ErrorEventCapacityPersistsTempUnsched(t *testi body := []byte(`{"type":"error","error":{"code":"server_error","type":"server_error","message":"Selected model is at capacity. Please try a different model."}}`) start := time.Now() - svc.persistOpenAIWSRateLimitSignal(context.Background(), account, http.Header{}, body, "server_error", "server_error", "Selected model is at capacity. Please try a different model.") + svc.persistOpenAIWSRateLimitSignal(context.Background(), account, "gpt-5.1", http.Header{}, body, "server_error", "server_error", "Selected model is at capacity. Please try a different model.") require.Empty(t, repo.rateLimitCalls) require.Len(t, repo.tempCalls, 1) @@ -270,6 +270,28 @@ func TestOpenAIGatewayService_WSv2ErrorEventCapacityPersistsTempUnsched(t *testi require.Contains(t, repo.tempReasons[0], "openai_model_capacity") } +func TestOpenAIGatewayService_WSModelCapacityIsIsolatedAndSuccessClearsStreak(t *testing.T) { + t.Parallel() + + svc := &OpenAIGatewayService{} + account := &Account{ + ID: 505, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + } + body := []byte(`{"type":"response.failed","error":{"code":"model_capacity_exhausted","message":"Selected model is at capacity."}}`) + for range 2 { + require.True(t, svc.handleOpenAIWSModelCapacitySignal( + context.Background(), account, "gpt-5.1", http.StatusServiceUnavailable, nil, body, "Selected model is at capacity.", + )) + } + require.True(t, svc.isOpenAIAccountModelRuntimeBlocked(account, "gpt-5.1")) + require.False(t, svc.isOpenAIAccountModelRuntimeBlocked(account, "gpt-5.2"), "capacity cooldown must not block other models on the same account") + + svc.ReportOpenAIAccountScheduleResult(account.ID, true, nil, account.GetMappedModel("gpt-5.1")) + require.False(t, svc.isOpenAIAccountModelRuntimeBlocked(account, "gpt-5.1"), "successful WS turn must clear the same model streak") +} + func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_ErrorEventUsageLimitPersistsRateLimit(t *testing.T) { gin.SetMode(gin.TestMode) diff --git a/backend/internal/service/openai_ws_v2/passthrough_relay.go b/backend/internal/service/openai_ws_v2/passthrough_relay.go index b243c2c87..358f03e3e 100644 --- a/backend/internal/service/openai_ws_v2/passthrough_relay.go +++ b/backend/internal/service/openai_ws_v2/passthrough_relay.go @@ -4,6 +4,7 @@ import ( "context" "crypto/sha256" "encoding/hex" + "encoding/json" "errors" "io" "net" @@ -523,6 +524,23 @@ func runUpstreamToClient( observedEvent := observedUpstreamEvent{} switch msgType { case coderws.MessageText: + if !json.Valid(payload) { + invalidErr := errors.New("upstream websocket returned invalid JSON") + emitRelayTrace(onTrace, RelayTraceEvent{ + Stage: "invalid_upstream_json", + Direction: "upstream_to_client", + MessageType: relayMessageTypeString(msgType), + PayloadBytes: len(payload), + WroteDownstream: wroteDownstream, + Error: invalidErr.Error(), + }) + exitCh <- relayExitSignal{ + stage: "invalid_upstream_json", + err: invalidErr, + wroteDownstream: wroteDownstream, + } + return + } observedEvent = observeUpstreamMessage(state, payload, startAt, nowFn, onUsageParseFailure) case coderws.MessageBinary: // binary frame 直接透传，不进入 JSON 观测路径（避免无效解析开销）。 diff --git a/backend/internal/service/openai_ws_v2/passthrough_relay_test.go b/backend/internal/service/openai_ws_v2/passthrough_relay_test.go index ec464f739..6a60345c0 100644 --- a/backend/internal/service/openai_ws_v2/passthrough_relay_test.go +++ b/backend/internal/service/openai_ws_v2/passthrough_relay_test.go @@ -19,10 +19,11 @@ type passthroughTestFrame struct { } type passthroughTestFrameConn struct { - mu sync.Mutex - writes []passthroughTestFrame - readCh chan passthroughTestFrame - once sync.Once + mu sync.Mutex + writes []passthroughTestFrame + readCh chan passthroughTestFrame + once sync.Once + closeCalls atomic.Int32 } type delayedReadFrameConn struct { @@ -80,6 +81,7 @@ func (c *passthroughTestFrameConn) WriteFrame(ctx context.Context, msgType coder } func (c *passthroughTestFrameConn) Close() error { + c.closeCalls.Add(1) c.once.Do(func() { defer func() { _ = recover() }() close(c.readCh) @@ -87,6 +89,10 @@ func (c *passthroughTestFrameConn) Close() error { return nil } +func (c *passthroughTestFrameConn) CloseCalls() int32 { + return c.closeCalls.Load() +} + func (c *passthroughTestFrameConn) Writes() []passthroughTestFrame { c.mu.Lock() defer c.mu.Unlock() @@ -200,6 +206,31 @@ func TestRelay_BasicRelayAndUsage(t *testing.T) { require.JSONEq(t, `{"type":"response.completed","response":{"id":"resp_123","usage":{"input_tokens":7,"output_tokens":3,"input_tokens_details":{"cached_tokens":2,"cache_write_tokens":1}}}}`, string(clientWrites[0].payload)) } +func TestRelay_RejectsMalformedUpstreamJSONWithoutForwarding(t *testing.T) { + t.Parallel() + + clientConn := newPassthroughTestFrameConn(nil, false) + upstreamConn := newPassthroughTestFrameConn([]passthroughTestFrame{{ + msgType: coderws.MessageText, + payload: []byte(`{"type":"response.completed"`), + }}, true) + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + _, relayExit := Relay( + ctx, + clientConn, + upstreamConn, + []byte(`{"type":"response.create","model":"gpt-5.3-codex"}`), + RelayOptions{}, + ) + require.NotNil(t, relayExit) + require.Equal(t, "invalid_upstream_json", relayExit.Stage) + require.ErrorContains(t, relayExit.Err, "invalid JSON") + require.Empty(t, clientConn.Writes(), "invalid upstream JSON must never reach the client") + require.GreaterOrEqual(t, upstreamConn.CloseCalls(), int32(1), "bad upstream connection must be closed") +} + func TestRelay_CountsImageGenerationOutput(t *testing.T) { t.Parallel() diff --git a/backend/internal/service/openai_ws_v2_passthrough_adapter.go b/backend/internal/service/openai_ws_v2_passthrough_adapter.go index 248abbf18..b8e9ece80 100644 --- a/backend/internal/service/openai_ws_v2_passthrough_adapter.go +++ b/backend/internal/service/openai_ws_v2_passthrough_adapter.go @@ -2,6 +2,7 @@ package service import ( "context" + "encoding/json" "errors" "fmt" "net/http" @@ -47,6 +48,14 @@ func (c *openAIWSPolicyEnforcingFrameConn) ReadFrame(ctx context.Context) (coder if err != nil { return msgType, payload, err } + if msgType == coderws.MessageText && !json.Valid(payload) { + invalidErr := errors.New("invalid websocket request JSON") + return msgType, nil, NewOpenAIWSClientCloseError( + coderws.StatusPolicyViolation, + "invalid websocket request payload", + invalidErr, + ) + } if c.filter == nil { return msgType, payload, nil } @@ -63,6 +72,74 @@ func (c *openAIWSPolicyEnforcingFrameConn) ReadFrame(ctx context.Context) (coder return msgType, updated, nil } +// openAIWSPassthroughTurnLifecycle serializes response.create turns on a +// passthrough connection. The ingress hooks own concurrency slots, so a new +// turn must not acquire its slots until the preceding terminal/error callback +// has released them. +type openAIWSPassthroughTurnLifecycle struct { + mu sync.Mutex + hooks *OpenAIWSIngressHooks + nextTurn int + activeTurn int +} + +func newOpenAIWSPassthroughTurnLifecycle(hooks *OpenAIWSIngressHooks) *openAIWSPassthroughTurnLifecycle { + return &openAIWSPassthroughTurnLifecycle{hooks: hooks, nextTurn: 1} +} + +func (l *openAIWSPassthroughTurnLifecycle) begin() (int, error) { + if l == nil { + return 0, errors.New("passthrough turn lifecycle is unavailable") + } + l.mu.Lock() + defer l.mu.Unlock() + if l.activeTurn > 0 { + return 0, NewOpenAIWSClientCloseError( + coderws.StatusPolicyViolation, + "parallel response.create turns are not supported", + nil, + ) + } + turn := l.nextTurn + if turn <= 0 { + turn = 1 + } + if l.hooks != nil && l.hooks.BeforeTurn != nil { + if err := l.hooks.BeforeTurn(turn); err != nil { + return 0, err + } + } + l.activeTurn = turn + l.nextTurn = turn + 1 + return turn, nil +} + +func (l *openAIWSPassthroughTurnLifecycle) finish(result *OpenAIForwardResult, turnErr error) (int, bool) { + if l == nil { + return 0, false + } + l.mu.Lock() + defer l.mu.Unlock() + turn := l.activeTurn + if turn <= 0 { + return 0, false + } + l.activeTurn = 0 + if l.hooks != nil && l.hooks.AfterTurn != nil { + l.hooks.AfterTurn(turn, result, turnErr) + } + return turn, true +} + +func (l *openAIWSPassthroughTurnLifecycle) hasActive() bool { + if l == nil { + return false + } + l.mu.Lock() + defer l.mu.Unlock() + return l.activeTurn > 0 +} + func (c *openAIWSPolicyEnforcingFrameConn) WriteFrame(ctx context.Context, msgType coderws.MessageType, payload []byte) error { if c == nil || c.inner == nil { return errOpenAIWSConnClosed @@ -170,6 +247,11 @@ func (c *openAIWSClientFrameConn) WriteFrame(ctx context.Context, msgType coderw if ctx == nil { ctx = context.Background() } + if msgType == coderws.MessageText { + if normalized, changed := normalizeCompletedImageGenerationStatus(payload); changed { + payload = normalized + } + } return c.conn.Write(ctx, msgType, payload) } @@ -236,6 +318,10 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough( // silently passed through, defeating the policy on every frame after // the first. capturedSessionModel := openAIWSPassthroughPolicyModelForFrame(account, firstClientMessage) + imageGenerationAllowed := GroupAllowsImageGeneration(apiKeyGroup(getAPIKeyFromContext(c))) + if permissionErr := openAIWSImageGenerationPermissionError(imageGenerationAllowed, capturedSessionModel, firstClientMessage); permissionErr != nil { + return permissionErr + } updatedFirst, blocked, policyErr := s.applyOpenAIFastPolicyToWSResponseCreate(ctx, account, capturedSessionModel, firstClientMessage) if policyErr != nil { return fmt.Errorf("apply openai fast policy on first ws frame: %w", policyErr) @@ -380,6 +466,7 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough( } completedTurns := atomic.Int32{} + turnLifecycle := newOpenAIWSPassthroughTurnLifecycle(hooks) policyClientConn := &openAIWSPolicyEnforcingFrameConn{ inner: &openAIWSClientFrameConn{conn: clientConn}, // 注意线程安全：filter 仅在 runClientToUpstream 这一条 @@ -417,6 +504,9 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough( if model == "" { model = capturedSessionModel } + if permissionErr := openAIWSImageGenerationPermissionError(imageGenerationAllowed, model, payload); permissionErr != nil { + return payload, nil, permissionErr + } out, blocked, policyErr := s.applyOpenAIFastPolicyToWSResponseCreate(ctx, account, model, payload) if policyErr == nil && blocked == nil && strings.TrimSpace(gjson.GetBytes(out, "type").String()) == "response.create" { @@ -426,29 +516,6 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough( } out = cleanedOut } - // 多轮 passthrough billing：仅在成功（non-block / non-err） - // 的 response.create 帧上更新 requestServiceTierPtr，使用 - // filter 处理后的 payload，与首帧 policy-after-extract 语义 - // 保持一致（参见上方 extractOpenAIServiceTierFromBody 注释）。 - // - 非 response.create 帧（response.cancel / - // conversation.item.create / session.update 等）不携带 - // per-response service_tier，不应覆盖前一轮值。 - // - blocked != nil：该帧不会发送上游，billing tier 应保持 - // 上一轮值。 - // - policyErr != nil：异常路径，保持上一轮值。 - // - 不带 service_tier 的 response.create 会让 - // extractOpenAIServiceTierFromBody 返回 nil；这里有意 - // 覆盖（Store(nil)），因为 OpenAI 上游对该帧实际不传 - // service_tier 时按 default 处理，billing 应如实反映。 - if policyErr == nil && blocked == nil && - strings.TrimSpace(gjson.GetBytes(out, "type").String()) == "response.create" { - requestServiceTierPtr.Store(extractOpenAIServiceTierFromBody(out)) - frameModel := strings.TrimSpace(gjson.GetBytes(out, "model").String()) - if frameModel == "" { - frameModel = requestModel - } - imageBillingConfigStore.Store(resolveOpenAIResponseImageBillingConfigFromBody(openAIResponsesEndpoint, frameModel, out)) - } return out, blocked, policyErr }, onBlock: func(blocked *OpenAIFastBlockedError) { @@ -480,8 +547,25 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough( truncateOpenAIWSLogValue(usageRaw, openAIWSLogValueMaxLen), ) }, + BeforeClientFrame: func(msgType coderws.MessageType, payload []byte) error { + if msgType != coderws.MessageText || strings.TrimSpace(gjson.GetBytes(payload, "type").String()) != "response.create" { + return nil + } + if _, beginErr := turnLifecycle.begin(); beginErr != nil { + return beginErr + } + // Update per-turn accounting only after lifecycle acquisition. A + // rejected pipelined response.create must not overwrite the active + // turn's service tier or image billing configuration. + requestServiceTierPtr.Store(extractOpenAIServiceTierFromBody(payload)) + frameModel := strings.TrimSpace(gjson.GetBytes(payload, "model").String()) + if frameModel == "" { + frameModel = requestModel + } + imageBillingConfigStore.Store(resolveOpenAIResponseImageBillingConfigFromBody(openAIResponsesEndpoint, frameModel, payload)) + return nil + }, OnTurnComplete: func(turn openaiwsv2.RelayTurnResult) { - turnNo := int(completedTurns.Add(1)) turnResult := &OpenAIForwardResult{ RequestID: turn.RequestID, Usage: OpenAIUsage{ @@ -498,6 +582,7 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough( ImageCount: turn.Usage.ImageCount, }, Model: turn.RequestModel, + UpstreamModel: turn.RequestModel, ServiceTier: requestServiceTierPtr.Load(), Stream: true, OpenAIWSMode: true, @@ -506,6 +591,17 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough( FirstTokenMs: turn.FirstTokenMs, } applyOpenAIResponseImageAccounting(turnResult, imageBillingConfigStore.Load()) + turnNo, finished := turnLifecycle.finish(turnResult, nil) + if !finished { + logOpenAIWSV2Passthrough( + "relay_terminal_without_active_turn account_id=%d request_id=%s terminal_event=%s", + account.ID, + truncateOpenAIWSLogValue(turnResult.RequestID, openAIWSIDValueMaxLen), + truncateOpenAIWSLogValue(turn.TerminalEventType, openAIWSLogValueMaxLen), + ) + return + } + completedTurns.Add(1) 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", account.ID, @@ -518,9 +614,6 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough( turnResult.Usage.OutputTokens, turnResult.Usage.CacheReadInputTokens, ) - if hooks != nil && hooks.AfterTurn != nil { - hooks.AfterTurn(turnNo, turnResult, nil) - } }, OnTrace: func(event openaiwsv2.RelayTraceEvent) { logOpenAIWSV2Passthrough( @@ -554,6 +647,7 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough( ImageCount: relayResult.Usage.ImageCount, }, Model: relayResult.RequestModel, + UpstreamModel: relayResult.RequestModel, ServiceTier: requestServiceTierPtr.Load(), Stream: true, OpenAIWSMode: true, @@ -565,6 +659,15 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough( turnCount := int(completedTurns.Load()) if relayExit == nil { + if turnLifecycle.hasActive() { + turnErr := wrapOpenAIWSIngressTurnError( + "incomplete_turn", + errors.New("upstream websocket closed before a terminal response event"), + relayResult.UpstreamToClientFrames > 0, + ) + turnLifecycle.finish(nil, turnErr) + return turnErr + } logOpenAIWSV2Passthrough( "relay_completed account_id=%d request_id=%s terminal_event=%s duration_ms=%d c2u_frames=%d u2c_frames=%d dropped_frames=%d turns=%d", account.ID, @@ -576,10 +679,6 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough( relayResult.DroppedDownstreamFrames, turnCount, ) - // 正常路径按 terminal 事件逐 turn 已回调；仅在零 turn 场景兜底回调一次。 - if turnCount == 0 && hooks != nil && hooks.AfterTurn != nil { - hooks.AfterTurn(1, result, nil) - } return nil } logOpenAIWSV2Passthrough( @@ -608,9 +707,7 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough( relayErr, relayExit.WroteDownstream, ) - if hooks != nil && hooks.AfterTurn != nil { - hooks.AfterTurn(turnCount+1, nil, turnErr) - } + turnLifecycle.finish(nil, turnErr) return turnErr } diff --git a/backend/internal/service/ops_cleanup_service.go b/backend/internal/service/ops_cleanup_service.go index 2dba11b5f..c1cbce03b 100644 --- a/backend/internal/service/ops_cleanup_service.go +++ b/backend/internal/service/ops_cleanup_service.go @@ -3,6 +3,8 @@ package service import ( "context" "database/sql" + "encoding/json" + "errors" "fmt" "io" "strings" @@ -22,6 +24,8 @@ const ( opsCleanupLeaderLockKeyDefault = "ops:cleanup:leader" opsCleanupLeaderLockTTLDefault = 30 * time.Minute + opsCleanupLeaderLockTTLGrace = 5 * time.Minute + opsCleanupReconcileInterval = time.Minute ) var opsCleanupCronParser = cron.NewParser(cron.Minute | cron.Hour | cron.Dom | cron.Month | cron.Dow) @@ -43,6 +47,7 @@ return 0 // 统一共享 cron schedule + leader lock + heartbeat，避免再引一套调度。 type OpsCleanupService struct { opsRepo OpsRepository + settingRepo SettingRepository db *sql.DB redisClient *redis.Client cfg *config.Config @@ -51,7 +56,18 @@ type OpsCleanupService struct { instanceID string - cron *cron.Cron + cronMu sync.Mutex + cron *cron.Cron + cronEntryID cron.EntryID + stopped bool + + scheduleStateInitialized bool + appliedSchedule string + appliedEnabled bool + + lifecycleCtx context.Context + lifecycleCancel context.CancelFunc + reconcileWG sync.WaitGroup startOnce sync.Once stopOnce sync.Once @@ -65,6 +81,7 @@ type opsCleanupArchiveCreator interface { func NewOpsCleanupService( opsRepo OpsRepository, + settingRepo SettingRepository, db *sql.DB, redisClient *redis.Client, cfg *config.Config, @@ -73,6 +90,7 @@ func NewOpsCleanupService( ) *OpsCleanupService { return &OpsCleanupService{ opsRepo: opsRepo, + settingRepo: settingRepo, db: db, redisClient: redisClient, cfg: cfg, @@ -86,40 +104,47 @@ func (s *OpsCleanupService) Start() { if s == nil { return } - if s.cfg != nil && !s.cfg.Ops.Enabled { + if s.cfg == nil { + logger.LegacyPrintf("service.ops_cleanup", "[OpsCleanup] not started (missing config)") return } - if s.cfg != nil && !s.cfg.Ops.Cleanup.Enabled { - logger.LegacyPrintf("service.ops_cleanup", "[OpsCleanup] not started (disabled)") + if !s.cfg.Ops.Enabled { return } if s.opsRepo == nil || s.db == nil { logger.LegacyPrintf("service.ops_cleanup", "[OpsCleanup] not started (missing deps)") return } + if err := s.cfg.Ops.Cleanup.Validate(); err != nil { + logger.LegacyPrintf("service.ops_cleanup", "[OpsCleanup] not started (invalid config): %v", err) + return + } + loc, err := loadOpsCleanupLocation(s.cfg.Timezone) + if err != nil { + logger.LegacyPrintf("service.ops_cleanup", "[OpsCleanup] not started (invalid config): %v", err) + return + } s.startOnce.Do(func() { - schedule := "0 2 * * *" - if s.cfg != nil && strings.TrimSpace(s.cfg.Ops.Cleanup.Schedule) != "" { - schedule = strings.TrimSpace(s.cfg.Ops.Cleanup.Schedule) - } - - loc := time.Local - if s.cfg != nil && strings.TrimSpace(s.cfg.Timezone) != "" { - if parsed, err := time.LoadLocation(strings.TrimSpace(s.cfg.Timezone)); err == nil && parsed != nil { - loc = parsed - } - } - c := cron.New(cron.WithParser(opsCleanupCronParser), cron.WithLocation(loc)) - _, err := c.AddFunc(schedule, func() { s.runScheduled() }) - if err != nil { - logger.LegacyPrintf("service.ops_cleanup", "[OpsCleanup] not started (invalid schedule=%q): %v", schedule, err) - return - } + lifecycleCtx, lifecycleCancel := context.WithCancel(context.Background()) + s.cronMu.Lock() s.cron = c - s.cron.Start() - logger.LegacyPrintf("service.ops_cleanup", "[OpsCleanup] started (schedule=%q tz=%s)", schedule, loc.String()) + s.stopped = false + s.lifecycleCtx = lifecycleCtx + s.lifecycleCancel = lifecycleCancel + s.cronMu.Unlock() + c.Start() + + ctx, cancel := context.WithTimeout(lifecycleCtx, 5*time.Second) + loadErr := s.ReconcileDataRetentionSettings(ctx) + cancel() + if loadErr != nil { + logger.LegacyPrintf("service.ops_cleanup", "[OpsCleanup] started without a cleanup entry (load dynamic settings failed): %v", loadErr) + } + + s.reconcileWG.Add(1) + go s.runSettingsReconcileLoop(lifecycleCtx) }) } @@ -128,23 +153,89 @@ func (s *OpsCleanupService) Stop() { return } s.stopOnce.Do(func() { - if s.cron != nil { - ctx := s.cron.Stop() + s.cronMu.Lock() + c := s.cron + cancel := s.lifecycleCancel + s.cron = nil + s.cronEntryID = 0 + s.stopped = true + s.cronMu.Unlock() + if cancel != nil { + cancel() + } + if c != nil { + ctx := c.Stop() select { case <-ctx.Done(): case <-time.After(3 * time.Second): logger.LegacyPrintf("service.ops_cleanup", "[OpsCleanup] cron stop timed out") } } + reconcileDone := make(chan struct{}) + go func() { + s.reconcileWG.Wait() + close(reconcileDone) + }() + select { + case <-reconcileDone: + case <-time.After(3 * time.Second): + logger.LegacyPrintf("service.ops_cleanup", "[OpsCleanup] settings reconcile stop timed out") + } }) } -func (s *OpsCleanupService) runScheduled() { - if s == nil || s.db == nil || s.opsRepo == nil { +func (s *OpsCleanupService) runSettingsReconcileLoop(ctx context.Context) { + defer s.reconcileWG.Done() + ticker := time.NewTicker(opsCleanupReconcileInterval) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + reconcileCtx, cancel := context.WithTimeout(ctx, 5*time.Second) + err := s.ReconcileDataRetentionSettings(reconcileCtx) + cancel() + if err != nil && !errors.Is(err, context.Canceled) { + logger.LegacyPrintf("service.ops_cleanup", "[OpsCleanup] reconcile dynamic settings failed: %v", err) + } + } + } +} + +func (s *OpsCleanupService) runScheduled(triggerSchedule string) { + if s == nil || s.db == nil || s.opsRepo == nil || s.cfg == nil { + return + } + parentCtx := s.lifecycleCtx + if parentCtx == nil { + parentCtx = context.Background() + } + settingsCtx, settingsCancel := context.WithTimeout(parentCtx, 5*time.Second) + cleanupCfg, err := s.loadEffectiveCleanupConfig(settingsCtx) + settingsCancel() + if err != nil { + logger.LegacyPrintf("service.ops_cleanup", "[OpsCleanup] cleanup rejected (invalid dynamic settings): %v", err) + return + } + if !cleanupCfg.Enabled { + return + } + if err := s.applyCleanupSchedule(cleanupCfg); err != nil { + logger.LegacyPrintf("service.ops_cleanup", "[OpsCleanup] cleanup rejected (schedule reconcile failed): %v", err) + return + } + if !opsCleanupTriggerMatches(triggerSchedule, cleanupCfg.Schedule) { + logger.LegacyPrintf("service.ops_cleanup", "[OpsCleanup] skipped stale cron callback (effective schedule=%q)", cleanupCfg.Schedule) + return + } + loc, err := loadOpsCleanupLocation(s.cfg.Timezone) + if err != nil { + logger.LegacyPrintf("service.ops_cleanup", "[OpsCleanup] cleanup rejected (invalid timezone): %v", err) return } - ctx, cancel := context.WithTimeout(context.Background(), 30*time.Minute) + ctx, cancel := context.WithTimeout(parentCtx, time.Duration(cleanupCfg.RunTimeoutSeconds)*time.Second) defer cancel() release, ok := s.tryAcquireLeaderLock(ctx) @@ -158,7 +249,7 @@ func (s *OpsCleanupService) runScheduled() { startedAt := time.Now().UTC() runAt := startedAt - counts, err := s.runCleanupOnce(ctx) + counts, err := s.runCleanupOnceWithConfig(ctx, cleanupCfg, loc) if err != nil { s.recordHeartbeatError(runAt, time.Since(startedAt), err) logger.LegacyPrintf("service.ops_cleanup", "[OpsCleanup] cleanup failed: %v", err) @@ -179,6 +270,13 @@ type opsCleanupDeletedCounts struct { dailyPreagg int64 } +type opsCleanupWindow struct { + Start time.Time + End time.Time +} + +type opsCleanupWindowExporter func(context.Context, opsCleanupWindow) (io.ReadCloser, error) + func (c opsCleanupDeletedCounts) String() string { return fmt.Sprintf( "error_logs=%d retry_attempts=%d alert_events=%d system_logs=%d log_audits=%d system_metrics=%d hourly_preagg=%d daily_preagg=%d", @@ -193,105 +291,246 @@ func (c opsCleanupDeletedCounts) String() string { ) } -// opsCleanupPlan 把"保留天数"翻译成具体的清理动作。 -// - days < 0 → 跳过该项清理（ok=false），保留兼容老数据 -// - days == 0 → TRUNCATE TABLE（O(1) 全清），truncate=true -// - days > 0 → 批量 DELETE 早于 now-N天 的行，cutoff = now - N 天 -// -// 之所以 days==0 走 TRUNCATE 而非"now+24h cutoff + DELETE"： -// - 速度从 O(N) 降到 O(1)，对百万行级表毫秒完成 -// - 无 WAL 写入、无后续 VACUUM 压力 -// - 这些 ops 表只有 cleanup 任务自己写，TRUNCATE 的 ACCESS EXCLUSIVE 锁影响可忽略 -func opsCleanupPlan(now time.Time, days int) (cutoff time.Time, truncate, ok bool) { - if days < 0 { - return time.Time{}, false, false +// opsCleanupPlan 把"保留天数"翻译成清理截止时间；0 明确表示禁用该目标。 +func opsCleanupPlan(now time.Time, days int) (cutoff time.Time, ok bool) { + if days <= 0 { + return time.Time{}, false } - if days == 0 { - return time.Time{}, true, true - } - return now.AddDate(0, 0, -days), false, true + dayStart := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location()) + return dayStart.AddDate(0, 0, -days), true } -func (s *OpsCleanupService) runCleanupOnce(ctx context.Context) (opsCleanupDeletedCounts, error) { - out := opsCleanupDeletedCounts{} - if s == nil || s.db == nil || s.cfg == nil { - return out, nil +func loadOpsCleanupLocation(value string) (*time.Location, error) { + name := strings.TrimSpace(value) + if name == "" { + return nil, fmt.Errorf("ops cleanup timezone is required") + } + loc, err := time.LoadLocation(name) + if err != nil { + return nil, fmt.Errorf("load ops cleanup timezone %q: %w", name, err) } + return loc, nil +} - batchSize := 5000 +func opsCleanupTriggerMatches(triggerSchedule, effectiveSchedule string) bool { + return strings.TrimSpace(triggerSchedule) == strings.TrimSpace(effectiveSchedule) +} - now := time.Now().UTC() +func mergeOpsCleanupDataRetention(base config.OpsCleanupConfig, retention OpsDataRetentionSettings) (config.OpsCleanupConfig, error) { + if err := validateOpsDataRetentionSettings(retention); err != nil { + return config.OpsCleanupConfig{}, err + } + base.Enabled = base.Enabled && retention.CleanupEnabled + base.Schedule = strings.TrimSpace(retention.CleanupSchedule) + base.ErrorLogRetentionDays = retention.ErrorLogRetentionDays + base.MinuteMetricsRetentionDays = retention.MinuteMetricsRetentionDays + base.HourlyMetricsRetentionDays = retention.HourlyMetricsRetentionDays + if err := base.Validate(); err != nil { + return config.OpsCleanupConfig{}, err + } + return base, nil +} - // runOne 把"truncate? cutoff? batched delete?"封装到一处， - // 让三组清理（错误日志类 / 分钟指标 / 小时+日预聚合）调用方只关心表名和列名。 - runOne := func(truncate bool, cutoff time.Time, table, timeCol string, castDate bool) (int64, error) { - if truncate { - return truncateOpsTable(ctx, s.db, table) +func (s *OpsCleanupService) loadEffectiveCleanupConfig(ctx context.Context) (config.OpsCleanupConfig, error) { + if s == nil || s.cfg == nil { + return config.OpsCleanupConfig{}, fmt.Errorf("ops cleanup service is not configured") + } + base := s.cfg.Ops.Cleanup + if err := base.Validate(); err != nil { + return config.OpsCleanupConfig{}, err + } + if !base.Enabled || s.settingRepo == nil { + return base, nil + } + if ctx == nil { + ctx = context.Background() + } + raw, err := s.settingRepo.GetValue(ctx, SettingKeyOpsAdvancedSettings) + if err != nil { + if errors.Is(err, ErrSettingNotFound) { + return base, nil } - return deleteOldRowsByID(ctx, s.db, table, timeCol, cutoff, batchSize, castDate) + return config.OpsCleanupConfig{}, fmt.Errorf("load ops advanced settings: %w", err) + } + advanced := defaultOpsAdvancedSettings() + if err := json.Unmarshal([]byte(raw), advanced); err != nil { + return config.OpsCleanupConfig{}, fmt.Errorf("decode ops advanced settings: %w", err) + } + if err := validateOpsDataRetentionSettings(advanced.DataRetention); err != nil { + return config.OpsCleanupConfig{}, fmt.Errorf("validate ops advanced data retention: %w", err) + } + normalizeOpsAdvancedSettings(advanced) + return mergeOpsCleanupDataRetention(base, advanced.DataRetention) +} + +// ReconcileDataRetentionSettings treats the persisted settings row as the only source of truth. +// It is safe to call concurrently and is run periodically by every application instance. +func (s *OpsCleanupService) ReconcileDataRetentionSettings(ctx context.Context) error { + if s == nil || s.cfg == nil { + return fmt.Errorf("ops cleanup service is not configured") + } + if !s.cfg.Ops.Enabled { + return nil + } + cleanupCfg, err := s.loadEffectiveCleanupConfig(ctx) + if err != nil { + return err } + return s.applyCleanupSchedule(cleanupCfg) +} - // Error-like tables: error logs / retry attempts / alert events / system logs / cleanup audits. - if cutoff, truncate, ok := opsCleanupPlan(now, s.cfg.Ops.Cleanup.ErrorLogRetentionDays); ok { - if err := s.archiveOpsErrorLogsForCleanup(ctx, cutoff, truncate); err != nil { - return out, err +func (s *OpsCleanupService) applyCleanupSchedule(cleanupCfg config.OpsCleanupConfig) error { + if err := cleanupCfg.Validate(); err != nil { + return err + } + schedule := strings.TrimSpace(cleanupCfg.Schedule) + if cleanupCfg.Enabled { + if _, err := opsCleanupCronParser.Parse(schedule); err != nil { + return fmt.Errorf("invalid ops cleanup schedule %q: %w", schedule, err) } - n, err := runOne(truncate, cutoff, "ops_error_logs", "created_at", false) + } + + s.cronMu.Lock() + defer s.cronMu.Unlock() + if s.stopped { + return fmt.Errorf("ops cleanup scheduler is stopped") + } + if s.cron == nil { + return fmt.Errorf("ops cleanup scheduler is not initialized") + } + if s.scheduleStateInitialized && s.appliedEnabled == cleanupCfg.Enabled && s.appliedSchedule == schedule { + return nil + } + + var newEntryID cron.EntryID + if cleanupCfg.Enabled { + triggerSchedule := schedule + entryID, err := s.cron.AddFunc(schedule, func() { s.runScheduled(triggerSchedule) }) if err != nil { - return out, err + return fmt.Errorf("add ops cleanup schedule %q: %w", schedule, err) } - out.errorLogs = n + newEntryID = entryID + } + oldEntryID := s.cronEntryID + s.cronEntryID = newEntryID + s.scheduleStateInitialized = true + s.appliedEnabled = cleanupCfg.Enabled + s.appliedSchedule = schedule + if oldEntryID != 0 { + s.cron.Remove(oldEntryID) + } + if cleanupCfg.Enabled { + logger.LegacyPrintf("service.ops_cleanup", "[OpsCleanup] schedule applied (schedule=%q)", schedule) + } else { + logger.LegacyPrintf("service.ops_cleanup", "[OpsCleanup] schedule disabled by advanced settings") + } + return nil +} + +func (s *OpsCleanupService) runCleanupOnce(ctx context.Context) (opsCleanupDeletedCounts, error) { + out := opsCleanupDeletedCounts{} + if s == nil || s.db == nil || s.cfg == nil { + return out, fmt.Errorf("ops cleanup service is not configured") + } + cleanupCfg, err := s.loadEffectiveCleanupConfig(ctx) + if err != nil { + return out, err + } + if !cleanupCfg.Enabled { + return out, nil + } + loc, err := loadOpsCleanupLocation(s.cfg.Timezone) + if err != nil { + return out, err + } + return s.runCleanupOnceWithConfig(ctx, cleanupCfg, loc) +} - n, err = runOne(truncate, cutoff, "ops_retry_attempts", "created_at", false) +func (s *OpsCleanupService) runCleanupOnceWithConfig( + ctx context.Context, + cleanupCfg config.OpsCleanupConfig, + loc *time.Location, +) (opsCleanupDeletedCounts, error) { + out := opsCleanupDeletedCounts{} + if s == nil || s.db == nil || loc == nil { + return out, fmt.Errorf("ops cleanup service is not configured") + } + if err := cleanupCfg.Validate(); err != nil { + return out, err + } + if !cleanupCfg.Enabled { + return out, nil + } + now := time.Now().In(loc) + + runDelete := func(cutoff time.Time, table, timeCol string, castDate bool) (int64, error) { + deleteCtx, cancel := context.WithTimeout(ctx, time.Duration(cleanupCfg.DeleteTimeoutSeconds)*time.Second) + defer cancel() + return deleteOldRowsByID(deleteCtx, s.db, table, timeCol, cutoff, cleanupCfg.DeleteBatchSize, castDate) + } + var cleanupErr error + recordError := func(scope string, err error) { if err != nil { - return out, err + cleanupErr = errors.Join(cleanupErr, fmt.Errorf("%s: %w", scope, err)) } - out.retryAttempts = n + } - n, err = runOne(truncate, cutoff, "ops_alert_events", "created_at", false) - if err != nil { - return out, err + // Archive-backed log tables run first so auxiliary-table deletion cannot consume + // the run budget before both durable archives have advanced. + if cutoff, ok := opsCleanupPlan(now, cleanupCfg.ErrorLogRetentionDays); ok { + n, err := s.cleanupOpsLogWindowsWithConfig(ctx, cleanupCfg, loc, "ops_error_logs", "created_at", cutoff, s.exportOpsErrorLogWindow) + if err == nil { + out.errorLogs = n } - out.alertEvents = n + recordError("cleanup ops_error_logs", err) - if err := s.archiveOpsSystemLogsForCleanup(ctx, cutoff, truncate); err != nil { - return out, err + n, err = s.cleanupOpsLogWindowsWithConfig(ctx, cleanupCfg, loc, "ops_system_logs", "created_at", cutoff, s.exportOpsSystemLogWindow) + if err == nil { + out.systemLogs = n } - n, err = runOne(truncate, cutoff, "ops_system_logs", "created_at", false) - if err != nil { - return out, err + recordError("cleanup ops_system_logs", err) + + n, err = runDelete(cutoff, "ops_retry_attempts", "created_at", false) + if err == nil { + out.retryAttempts = n } - out.systemLogs = n + recordError("cleanup ops_retry_attempts", err) - n, err = runOne(truncate, cutoff, "ops_system_log_cleanup_audits", "created_at", false) - if err != nil { - return out, err + n, err = runDelete(cutoff, "ops_alert_events", "created_at", false) + if err == nil { + out.alertEvents = n + } + recordError("cleanup ops_alert_events", err) + + n, err = runDelete(cutoff, "ops_system_log_cleanup_audits", "created_at", false) + if err == nil { + out.logAudits = n } - out.logAudits = n + recordError("cleanup ops_system_log_cleanup_audits", err) } // Minute-level metrics snapshots. - if cutoff, truncate, ok := opsCleanupPlan(now, s.cfg.Ops.Cleanup.MinuteMetricsRetentionDays); ok { - n, err := runOne(truncate, cutoff, "ops_system_metrics", "created_at", false) - if err != nil { - return out, err + if cutoff, ok := opsCleanupPlan(now, cleanupCfg.MinuteMetricsRetentionDays); ok { + n, err := runDelete(cutoff, "ops_system_metrics", "created_at", false) + if err == nil { + out.systemMetrics = n } - out.systemMetrics = n + recordError("cleanup ops_system_metrics", err) } // Pre-aggregation tables (hourly/daily). - if cutoff, truncate, ok := opsCleanupPlan(now, s.cfg.Ops.Cleanup.HourlyMetricsRetentionDays); ok { - n, err := runOne(truncate, cutoff, "ops_metrics_hourly", "bucket_start", false) - if err != nil { - return out, err + if cutoff, ok := opsCleanupPlan(now, cleanupCfg.HourlyMetricsRetentionDays); ok { + n, err := runDelete(cutoff, "ops_metrics_hourly", "bucket_start", false) + if err == nil { + out.hourlyPreagg = n } - out.hourlyPreagg = n + recordError("cleanup ops_metrics_hourly", err) - n, err = runOne(truncate, cutoff, "ops_metrics_daily", "bucket_date", true) - if err != nil { - return out, err + n, err = runDelete(cutoff, "ops_metrics_daily", "bucket_date", true) + if err == nil { + out.dailyPreagg = n } - out.dailyPreagg = n + recordError("cleanup ops_metrics_daily", err) } // Channel monitor 每日维护（聚合昨日明细 + 软删过期明细/聚合）。 @@ -303,76 +542,104 @@ func (s *OpsCleanupService) runCleanupOnce(ctx context.Context) (opsCleanupDelet } } - return out, nil + return out, cleanupErr } -func (s *OpsCleanupService) archiveOpsErrorLogsForCleanup(ctx context.Context, cutoff time.Time, truncate bool) error { - if s == nil || s.db == nil || s.opsRepo == nil { - return nil +func (s *OpsCleanupService) exportOpsErrorLogWindow(ctx context.Context, window opsCleanupWindow) (io.ReadCloser, error) { + start, end := window.Start.UTC(), window.End.UTC() + return s.opsRepo.ExportErrorLogs(ctx, &OpsErrorLogCleanupFilter{StartTime: &start, EndTime: &end}) +} + +func (s *OpsCleanupService) exportOpsSystemLogWindow(ctx context.Context, window opsCleanupWindow) (io.ReadCloser, error) { + start, end := window.Start.UTC(), window.End.UTC() + return s.opsRepo.ExportSystemLogs(ctx, &OpsSystemLogCleanupFilter{StartTime: &start, EndTime: &end}) +} + +func (s *OpsCleanupService) cleanupOpsLogWindows( + ctx context.Context, + table string, + timeColumn string, + cutoff time.Time, + exportWindow opsCleanupWindowExporter, +) (int64, error) { + if s == nil || s.db == nil || s.cfg == nil { + return 0, fmt.Errorf("ops cleanup window dependencies are not configured") } - hasRows, err := hasRowsForOpsCleanup(ctx, s.db, "ops_error_logs", "created_at", cutoff, false, truncate) + cleanupCfg, err := s.loadEffectiveCleanupConfig(ctx) if err != nil { - return err - } - if !hasRows { - return nil + return 0, err } - if s.archiveCreator == nil { - return fmt.Errorf("ops cleanup archive creator is not configured") - } - - filter := &OpsErrorLogCleanupFilter{} - if !truncate { - end := cutoff.UTC() - filter.EndTime = &end + if !cleanupCfg.Enabled { + return 0, nil } - stream, err := s.opsRepo.ExportErrorLogs(ctx, filter) + loc, err := loadOpsCleanupLocation(s.cfg.Timezone) if err != nil { - if isMissingRelationError(err) { - return nil - } - return err + return 0, err } - return s.createOpsCleanupArchive(ctx, "ops_error_logs", cutoff, truncate, stream) + return s.cleanupOpsLogWindowsWithConfig(ctx, cleanupCfg, loc, table, timeColumn, cutoff, exportWindow) } -func (s *OpsCleanupService) archiveOpsSystemLogsForCleanup(ctx context.Context, cutoff time.Time, truncate bool) error { - if s == nil || s.db == nil || s.opsRepo == nil { - return nil - } - hasRows, err := hasRowsForOpsCleanup(ctx, s.db, "ops_system_logs", "created_at", cutoff, false, truncate) - if err != nil { - return err - } - if !hasRows { - return nil +func (s *OpsCleanupService) cleanupOpsLogWindowsWithConfig( + ctx context.Context, + cleanupCfg config.OpsCleanupConfig, + loc *time.Location, + table string, + timeColumn string, + cutoff time.Time, + exportWindow opsCleanupWindowExporter, +) (int64, error) { + if s == nil || s.db == nil || exportWindow == nil || loc == nil { + return 0, fmt.Errorf("ops cleanup window dependencies are not configured") } if s.archiveCreator == nil { - return fmt.Errorf("ops cleanup archive creator is not configured") + return 0, fmt.Errorf("ops cleanup archive creator is not configured") } - - filter := &OpsSystemLogCleanupFilter{} - if !truncate { - end := cutoff.UTC() - filter.EndTime = &end + if err := cleanupCfg.Validate(); err != nil { + return 0, err } - stream, err := s.opsRepo.ExportSystemLogs(ctx, filter) - if err != nil { - if isMissingRelationError(err) { - return nil + var deletedTotal int64 + for range cleanupCfg.MaxCatchupWindowsPerRun { + oldest, err := findOldestOpsCleanupTime(ctx, s.db, table, timeColumn, cutoff) + if err != nil { + return deletedTotal, err + } + if oldest == nil { + return deletedTotal, nil + } + window, err := buildOpsCleanupWindow(*oldest, cutoff, cleanupCfg.ArchiveWindowDays, loc) + if err != nil { + return deletedTotal, err + } + + archiveCtx, archiveCancel := context.WithTimeout(ctx, time.Duration(cleanupCfg.ArchiveTimeoutSeconds)*time.Second) + stream, err := exportWindow(archiveCtx, window) + if err == nil { + err = s.createOpsCleanupArchive(archiveCtx, table, window, stream) + } + archiveCancel() + if err != nil { + return deletedTotal, fmt.Errorf("archive %s window [%s,%s): %w", table, window.Start.Format(time.RFC3339), window.End.Format(time.RFC3339), err) + } + + deleteCtx, deleteCancel := context.WithTimeout(ctx, time.Duration(cleanupCfg.DeleteTimeoutSeconds)*time.Second) + deleted, err := deleteRowsByIDWindow(deleteCtx, s.db, table, timeColumn, window, cleanupCfg.DeleteBatchSize) + deleteCancel() + deletedTotal += deleted + if err != nil { + return deletedTotal, fmt.Errorf("delete %s window [%s,%s): %w", table, window.Start.Format(time.RFC3339), window.End.Format(time.RFC3339), err) } - return err } - return s.createOpsCleanupArchive(ctx, "ops_system_logs", cutoff, truncate, stream) + return deletedTotal, nil } -func (s *OpsCleanupService) createOpsCleanupArchive(ctx context.Context, table string, cutoff time.Time, truncate bool, stream io.ReadCloser) error { +func (s *OpsCleanupService) createOpsCleanupArchive(ctx context.Context, table string, window opsCleanupWindow, stream io.ReadCloser) error { if stream == nil { return fmt.Errorf("%s archive stream is nil", table) } + defer func() { _ = stream.Close() }() record, err := s.archiveCreator.CreateDataArchive(ctx, DataArchiveInput{ Stream: stream, - FileName: opsCleanupArchiveFileName(table, cutoff, truncate), + FileName: opsCleanupArchiveFileName(table, window, time.Now().UTC()), BackupType: table + "_archive", TriggeredBy: opsCleanupArchiveTriggeredBy, ExpireDays: s.opsArchiveExpireDays(), @@ -394,45 +661,46 @@ func (s *OpsCleanupService) opsArchiveExpireDays() int { return s.cfg.Ops.Cleanup.ArchiveExpireDays } -func opsCleanupArchiveFileName(table string, cutoff time.Time, truncate bool) string { - if truncate || cutoff.IsZero() { - return fmt.Sprintf("%s_full_%s.ndjson.gz", table, time.Now().UTC().Format("20060102_150405")) - } - return fmt.Sprintf("%s_before_%s.ndjson.gz", table, cutoff.UTC().Format("20060102_150405")) +func opsCleanupArchiveFileName(table string, window opsCleanupWindow, attemptAt time.Time) string { + return fmt.Sprintf( + "%s_%s_%s_%s.ndjson.gz", + table, + window.Start.UTC().Format("20060102_150405"), + window.End.UTC().Format("20060102_150405"), + attemptAt.UTC().Format("20060102_150405.000000000"), + ) } -func hasRowsForOpsCleanup( - ctx context.Context, - db *sql.DB, - table string, - timeColumn string, - cutoff time.Time, - castCutoffToDate bool, - truncate bool, -) (bool, error) { +func findOldestOpsCleanupTime(ctx context.Context, db *sql.DB, table, timeColumn string, cutoff time.Time) (*time.Time, error) { if db == nil { - return false, nil + return nil, fmt.Errorf("ops cleanup database is not configured") } - var query string - var args []any - if truncate { - query = fmt.Sprintf("SELECT EXISTS(SELECT 1 FROM %s LIMIT 1)", table) - } else { - where := fmt.Sprintf("%s < $1", timeColumn) - if castCutoffToDate { - where = fmt.Sprintf("%s < $1::date", timeColumn) - } - query = fmt.Sprintf("SELECT EXISTS(SELECT 1 FROM %s WHERE %s LIMIT 1)", table, where) - args = []any{cutoff} + var oldest sql.NullTime + query := fmt.Sprintf("SELECT MIN(%s) FROM %s WHERE %s < $1", timeColumn, table, timeColumn) + if err := db.QueryRowContext(ctx, query, cutoff.UTC()).Scan(&oldest); err != nil { + return nil, err } - var exists bool - if err := db.QueryRowContext(ctx, query, args...).Scan(&exists); err != nil { - if isMissingRelationError(err) { - return false, nil - } - return false, err + if !oldest.Valid { + return nil, nil + } + value := oldest.Time.UTC() + return &value, nil +} + +func buildOpsCleanupWindow(oldest, cutoff time.Time, windowDays int, loc *time.Location) (opsCleanupWindow, error) { + if windowDays <= 0 || loc == nil { + return opsCleanupWindow{}, fmt.Errorf("invalid ops cleanup window configuration") + } + localOldest := oldest.In(loc) + start := time.Date(localOldest.Year(), localOldest.Month(), localOldest.Day(), 0, 0, 0, 0, loc).UTC() + end := start.In(loc).AddDate(0, 0, windowDays).UTC() + if end.After(cutoff.UTC()) { + end = cutoff.UTC() + } + if !end.After(start) { + return opsCleanupWindow{}, fmt.Errorf("invalid ops cleanup window [%s,%s)", start.Format(time.RFC3339), end.Format(time.RFC3339)) } - return exists, nil + return opsCleanupWindow{Start: start, End: end}, nil } func deleteOldRowsByID( @@ -445,10 +713,10 @@ func deleteOldRowsByID( castCutoffToDate bool, ) (int64, error) { if db == nil { - return 0, nil + return 0, fmt.Errorf("ops cleanup database is not configured") } if batchSize <= 0 { - batchSize = 5000 + return 0, fmt.Errorf("ops cleanup delete batch size must be positive") } where := fmt.Sprintf("%s < $1", timeColumn) @@ -471,10 +739,6 @@ WHERE id IN (SELECT id FROM batch) for { res, err := db.ExecContext(ctx, q, cutoff, batchSize) if err != nil { - // If ops tables aren't present yet (partial deployments), treat as no-op. - if isMissingRelationError(err) { - return total, nil - } return total, err } affected, err := res.RowsAffected() @@ -482,51 +746,55 @@ WHERE id IN (SELECT id FROM batch) return total, err } total += affected - if affected == 0 { + if affected < int64(batchSize) { break } } return total, nil } -// truncateOpsTable 用 TRUNCATE TABLE 清空指定表，先 SELECT COUNT(*) 取得清空前行数用于 heartbeat。 -// -// 与 deleteOldRowsByID 的差异： -// - 不可指定 WHERE 条件，仅用于 days==0 的"清空全部"语义 -// - O(1) 释放表的物理存储页，毫秒级完成，无 WAL 写入、无 VACUUM 压力 -// - 需要 ACCESS EXCLUSIVE 锁，但 ops 表只有清理任务自己写入，瞬间锁影响可忽略 -// -// 表不存在（部分部署）静默返回 0，与 deleteOldRowsByID 保持一致。 -func truncateOpsTable(ctx context.Context, db *sql.DB, table string) (int64, error) { +func deleteRowsByIDWindow( + ctx context.Context, + db *sql.DB, + table string, + timeColumn string, + window opsCleanupWindow, + batchSize int, +) (int64, error) { if db == nil { - return 0, nil + return 0, fmt.Errorf("ops cleanup database is not configured") } - var count int64 - if err := db.QueryRowContext(ctx, fmt.Sprintf("SELECT COUNT(*) FROM %s", table)).Scan(&count); err != nil { - if isMissingRelationError(err) { - return 0, nil - } - return 0, fmt.Errorf("count %s: %w", table, err) + if batchSize <= 0 { + return 0, fmt.Errorf("ops cleanup delete batch size must be positive") } - if count == 0 { - return 0, nil + if !window.End.After(window.Start) { + return 0, fmt.Errorf("invalid ops cleanup delete window") } - if _, err := db.ExecContext(ctx, fmt.Sprintf("TRUNCATE TABLE %s", table)); err != nil { - if isMissingRelationError(err) { - return 0, nil + query := fmt.Sprintf(` +WITH batch AS ( + SELECT id FROM %s + WHERE %s >= $1 AND %s < $2 + ORDER BY id + LIMIT $3 +) +DELETE FROM %s +WHERE id IN (SELECT id FROM batch) +`, table, timeColumn, timeColumn, table) + var total int64 + for { + result, err := db.ExecContext(ctx, query, window.Start.UTC(), window.End.UTC(), batchSize) + if err != nil { + return total, err + } + affected, err := result.RowsAffected() + if err != nil { + return total, err + } + total += affected + if affected < int64(batchSize) { + return total, nil } - return 0, fmt.Errorf("truncate %s: %w", table, err) - } - return count, nil -} - -// isMissingRelationError 判断 PG 报错是否为"表不存在"，用于让清理任务在部分部署场景静默跳过。 -func isMissingRelationError(err error) bool { - if err == nil { - return false } - s := strings.ToLower(err.Error()) - return strings.Contains(s, "does not exist") && strings.Contains(s, "relation") } func (s *OpsCleanupService) tryAcquireLeaderLock(ctx context.Context) (func(), bool) { @@ -539,7 +807,7 @@ func (s *OpsCleanupService) tryAcquireLeaderLock(ctx context.Context) (func(), b } key := opsCleanupLeaderLockKeyDefault - ttl := opsCleanupLeaderLockTTLDefault + ttl := configuredOpsCleanupLeaderLockTTL(s.cfg) // Prefer Redis leader lock when available, but avoid stampeding the DB when Redis is flaky by // falling back to a DB advisory lock. @@ -550,7 +818,11 @@ func (s *OpsCleanupService) tryAcquireLeaderLock(ctx context.Context) (func(), b return nil, false } return func() { - _, _ = opsCleanupReleaseScript.Run(ctx, s.redisClient, []string{key}, s.instanceID).Result() + releaseCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + if _, releaseErr := opsCleanupReleaseScript.Run(releaseCtx, s.redisClient, []string{key}, s.instanceID).Result(); releaseErr != nil { + logger.LegacyPrintf("service.ops_cleanup", "[OpsCleanup] leader lock release failed: %v", releaseErr) + } }, true } // Redis error: fall back to DB advisory lock. @@ -570,6 +842,18 @@ func (s *OpsCleanupService) tryAcquireLeaderLock(ctx context.Context) (func(), b return release, true } +func configuredOpsCleanupLeaderLockTTL(cfg *config.Config) time.Duration { + ttl := opsCleanupLeaderLockTTLDefault + if cfg == nil || cfg.Ops.Cleanup.RunTimeoutSeconds <= 0 { + return ttl + } + runTTL := time.Duration(cfg.Ops.Cleanup.RunTimeoutSeconds)*time.Second + opsCleanupLeaderLockTTLGrace + if runTTL > ttl { + return runTTL + } + return ttl +} + func (s *OpsCleanupService) recordHeartbeatSuccess(runAt time.Time, duration time.Duration, counts opsCleanupDeletedCounts) { if s == nil || s.opsRepo == nil { return diff --git a/backend/internal/service/ops_cleanup_service_test.go b/backend/internal/service/ops_cleanup_service_test.go index 86657d272..912d05a5d 100644 --- a/backend/internal/service/ops_cleanup_service_test.go +++ b/backend/internal/service/ops_cleanup_service_test.go @@ -1,64 +1,506 @@ package service import ( + "context" + "encoding/json" + "errors" + "io" + "strings" "testing" "time" + + "github.com/DATA-DOG/go-sqlmock" + "github.com/Wei-Shaw/sub2api/internal/config" ) func TestOpsCleanupPlan(t *testing.T) { - now := time.Date(2026, 4, 29, 12, 0, 0, 0, time.UTC) + loc, err := time.LoadLocation("Asia/Shanghai") + if err != nil { + t.Fatal(err) + } + now := time.Date(2026, 4, 29, 12, 0, 0, 0, loc) + dayStart := time.Date(2026, 4, 29, 0, 0, 0, 0, loc) cases := []struct { - name string - days int - wantOK bool - wantTruncate bool - wantCutoff time.Time + name string + days int + wantOK bool + wantCutoff time.Time }{ {name: "negative skips", days: -1, wantOK: false}, - {name: "zero truncates", days: 0, wantOK: true, wantTruncate: true}, - {name: "positive yields past cutoff", days: 7, wantOK: true, wantCutoff: now.AddDate(0, 0, -7)}, + {name: "zero disables", days: 0, wantOK: false}, + {name: "positive yields natural-day cutoff", days: 7, wantOK: true, wantCutoff: dayStart.AddDate(0, 0, -7)}, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { - cutoff, truncate, ok := opsCleanupPlan(now, tc.days) + cutoff, ok := opsCleanupPlan(now, tc.days) if ok != tc.wantOK { t.Fatalf("ok = %v, want %v", ok, tc.wantOK) } if !ok { return } - if truncate != tc.wantTruncate { - t.Fatalf("truncate = %v, want %v", truncate, tc.wantTruncate) - } - if !tc.wantTruncate && !cutoff.Equal(tc.wantCutoff) { + if !cutoff.Equal(tc.wantCutoff) { t.Fatalf("cutoff = %v, want %v", cutoff, tc.wantCutoff) } }) } } -func TestIsMissingRelationError(t *testing.T) { - cases := []struct { - name string - err error - want bool - }{ - {name: "nil is not missing", err: nil, want: false}, - {name: "match relation does not exist", err: fakeErr(`pq: relation "ops_error_logs" does not exist`), want: true}, - {name: "match case-insensitive", err: fakeErr(`ERROR: Relation "x" Does Not Exist`), want: true}, - {name: "non-matching error", err: fakeErr("connection refused"), want: false}, +func TestLoadOpsCleanupLocation(t *testing.T) { + if _, err := loadOpsCleanupLocation(" "); err == nil { + t.Fatal("empty timezone must be rejected") } - for _, tc := range cases { - t.Run(tc.name, func(t *testing.T) { - if got := isMissingRelationError(tc.err); got != tc.want { - t.Fatalf("got %v, want %v", got, tc.want) - } - }) + if _, err := loadOpsCleanupLocation("Invalid/Timezone"); err == nil { + t.Fatal("invalid timezone must be rejected") + } + loc, err := loadOpsCleanupLocation(" Asia/Shanghai ") + if err != nil { + t.Fatal(err) + } + if loc.String() != "Asia/Shanghai" { + t.Fatalf("location = %q, want Asia/Shanghai", loc.String()) + } +} + +func TestConfiguredOpsCleanupLeaderLockTTL(t *testing.T) { + cfg := testOpsCleanupConfig() + if got := configuredOpsCleanupLeaderLockTTL(cfg); got != opsCleanupLeaderLockTTLDefault { + t.Fatalf("short run lock TTL = %s, want %s", got, opsCleanupLeaderLockTTLDefault) + } + cfg.Ops.Cleanup.RunTimeoutSeconds = 5 * 60 * 60 + want := 5*time.Hour + opsCleanupLeaderLockTTLGrace + if got := configuredOpsCleanupLeaderLockTTL(cfg); got != want { + t.Fatalf("long run lock TTL = %s, want %s", got, want) + } +} + +type opsCleanupArchiveStub struct { + create func(context.Context, DataArchiveInput) (*BackupRecord, error) +} + +func (s *opsCleanupArchiveStub) CreateDataArchive(ctx context.Context, input DataArchiveInput) (*BackupRecord, error) { + return s.create(ctx, input) +} + +func testOpsCleanupConfig() *config.Config { + return &config.Config{ + Timezone: "Asia/Shanghai", + Ops: config.OpsConfig{Enabled: true, Cleanup: config.OpsCleanupConfig{ + Enabled: true, + Schedule: "0 4 * * *", + ArchiveExpireDays: 30, + ArchiveWindowDays: 1, + MaxCatchupWindowsPerRun: 2, + ArchiveTimeoutSeconds: 1, + DeleteTimeoutSeconds: 1, + RunTimeoutSeconds: 10, + DeleteBatchSize: 5000, + ErrorLogRetentionDays: 30, + MinuteMetricsRetentionDays: 0, + HourlyMetricsRetentionDays: 0, + }}, + } +} + +func TestOpsCleanupEffectiveConfigUsesAdvancedSettings(t *testing.T) { + repo := newRuntimeSettingRepoStub() + advanced := defaultOpsAdvancedSettings() + advanced.DataRetention = OpsDataRetentionSettings{ + CleanupEnabled: true, + CleanupSchedule: "15 4 * * *", + ErrorLogRetentionDays: 11, + MinuteMetricsRetentionDays: 0, + HourlyMetricsRetentionDays: 22, + } + raw, err := json.Marshal(advanced) + if err != nil { + t.Fatal(err) + } + repo.values[SettingKeyOpsAdvancedSettings] = string(raw) + db, _, err := sqlmock.New() + if err != nil { + t.Fatal(err) + } + defer db.Close() + svc := NewOpsCleanupService(&opsRepoMock{}, repo, db, nil, testOpsCleanupConfig(), nil, nil) + + got, err := svc.loadEffectiveCleanupConfig(context.Background()) + if err != nil { + t.Fatal(err) + } + if !got.Enabled || got.Schedule != "15 4 * * *" { + t.Fatalf("effective enabled/schedule = %v/%q", got.Enabled, got.Schedule) + } + if got.ErrorLogRetentionDays != 11 || got.MinuteMetricsRetentionDays != 0 || got.HourlyMetricsRetentionDays != 22 { + t.Fatalf("effective retention = %d/%d/%d", got.ErrorLogRetentionDays, got.MinuteMetricsRetentionDays, got.HourlyMetricsRetentionDays) + } + if got.ArchiveExpireDays != 30 || got.DeleteBatchSize != 5000 { + t.Fatalf("static execution controls were not preserved: %+v", got) + } +} + +func TestOpsCleanupMalformedAdvancedSettingsFailsClosed(t *testing.T) { + repo := newRuntimeSettingRepoStub() + repo.values[SettingKeyOpsAdvancedSettings] = "{invalid" + db, _, err := sqlmock.New() + if err != nil { + t.Fatal(err) + } + defer db.Close() + svc := NewOpsCleanupService(&opsRepoMock{}, repo, db, nil, testOpsCleanupConfig(), nil, nil) + + if _, err := svc.loadEffectiveCleanupConfig(context.Background()); err == nil { + t.Fatal("malformed advanced settings must reject cleanup") + } +} + +func TestOpsCleanupReconcileDataRetentionReschedulesAndDisables(t *testing.T) { + db, _, err := sqlmock.New() + if err != nil { + t.Fatal(err) + } + defer db.Close() + repo := newRuntimeSettingRepoStub() + svc := NewOpsCleanupService(&opsRepoMock{}, repo, db, nil, testOpsCleanupConfig(), nil, nil) + svc.Start() + defer svc.Stop() + + advanced := defaultOpsAdvancedSettings() + advanced.DataRetention = OpsDataRetentionSettings{ + CleanupEnabled: true, + CleanupSchedule: "15 4 * * *", + ErrorLogRetentionDays: 30, + MinuteMetricsRetentionDays: 30, + HourlyMetricsRetentionDays: 30, + } + raw, err := json.Marshal(advanced) + if err != nil { + t.Fatal(err) + } + repo.values[SettingKeyOpsAdvancedSettings] = string(raw) + if err := svc.ReconcileDataRetentionSettings(context.Background()); err != nil { + t.Fatal(err) + } + svc.cronMu.Lock() + entryID := svc.cronEntryID + entries := svc.cron.Entries() + svc.cronMu.Unlock() + if entryID == 0 || len(entries) != 1 || entries[0].ID != entryID { + t.Fatalf("rescheduled entries = %+v, entryID=%d", entries, entryID) + } + if entries[0].Next.Hour() != 4 || entries[0].Next.Minute() != 15 { + t.Fatalf("next run = %s, want 04:15", entries[0].Next) + } + + advanced.DataRetention.CleanupEnabled = false + advanced.DataRetention.CleanupSchedule = "stale invalid cron" + raw, err = json.Marshal(advanced) + if err != nil { + t.Fatal(err) + } + repo.values[SettingKeyOpsAdvancedSettings] = string(raw) + if err := svc.ReconcileDataRetentionSettings(context.Background()); err != nil { + t.Fatal(err) + } + svc.cronMu.Lock() + defer svc.cronMu.Unlock() + if svc.cronEntryID != 0 || len(svc.cron.Entries()) != 0 { + t.Fatalf("disabled scheduler still has entries: %+v", svc.cron.Entries()) } } -type fakeErr string +func TestOpsCleanupTriggerMatchesEffectiveSchedule(t *testing.T) { + if !opsCleanupTriggerMatches(" 0 4 * * * ", "0 4 * * *") { + t.Fatal("equivalent cron specifications should match") + } + if opsCleanupTriggerMatches("0 4 * * *", "15 4 * * *") { + t.Fatal("stale cron specification must not match the effective schedule") + } +} -func (e fakeErr) Error() string { return string(e) } +func TestOpsCleanupStartupSettingsFailureCanReconcileLater(t *testing.T) { + repo := newRuntimeSettingRepoStub() + loadFails := true + repo.getValueFn = func(key string) (string, error) { + if loadFails { + return "", errors.New("temporary database error") + } + value, ok := repo.values[key] + if !ok { + return "", ErrSettingNotFound + } + return value, nil + } + db, _, err := sqlmock.New() + if err != nil { + t.Fatal(err) + } + defer db.Close() + svc := NewOpsCleanupService(&opsRepoMock{}, repo, db, nil, testOpsCleanupConfig(), nil, nil) + svc.Start() + defer svc.Stop() + svc.cronMu.Lock() + initialEntryID := svc.cronEntryID + svc.cronMu.Unlock() + if initialEntryID != 0 { + t.Fatalf("initial entry = %d, want none after settings failure", initialEntryID) + } + + advanced := defaultOpsAdvancedSettings() + raw, err := json.Marshal(advanced) + if err != nil { + t.Fatal(err) + } + repo.values[SettingKeyOpsAdvancedSettings] = string(raw) + loadFails = false + if err := svc.ReconcileDataRetentionSettings(context.Background()); err != nil { + t.Fatal(err) + } + svc.cronMu.Lock() + recoveredEntryID := svc.cronEntryID + lifecycleCtx := svc.lifecycleCtx + svc.cronMu.Unlock() + if recoveredEntryID == 0 { + t.Fatal("settings reconcile did not recover the cleanup entry") + } + svc.Stop() + select { + case <-lifecycleCtx.Done(): + default: + t.Fatal("Stop did not cancel the cleanup lifecycle context") + } +} + +func TestOpsCleanupErrorArchiveFailureStillAdvancesSystemLogs(t *testing.T) { + db, mock, err := sqlmock.New() + if err != nil { + t.Fatal(err) + } + defer db.Close() + oldest := time.Now().UTC().AddDate(0, 0, -40) + mock.ExpectQuery(`SELECT MIN\(created_at\) FROM ops_error_logs`). + WithArgs(sqlmock.AnyArg()). + WillReturnRows(sqlmock.NewRows([]string{"min"}).AddRow(oldest)) + mock.ExpectQuery(`SELECT MIN\(created_at\) FROM ops_system_logs`). + WithArgs(sqlmock.AnyArg()). + WillReturnRows(sqlmock.NewRows([]string{"min"}).AddRow(oldest)) + mock.ExpectExec(`DELETE FROM ops_system_logs`). + WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg(), 5000). + WillReturnResult(sqlmock.NewResult(0, 1)) + for _, table := range []string{"ops_retry_attempts", "ops_alert_events", "ops_system_log_cleanup_audits"} { + mock.ExpectExec(`DELETE FROM `+table). + WithArgs(sqlmock.AnyArg(), 5000). + WillReturnResult(sqlmock.NewResult(0, 0)) + } + repo := &opsRepoMock{ + ExportErrorLogsFn: func(context.Context, *OpsErrorLogCleanupFilter) (io.ReadCloser, error) { + return nil, errors.New("error archive unavailable") + }, + ExportSystemLogsFn: func(context.Context, *OpsSystemLogCleanupFilter) (io.ReadCloser, error) { + return io.NopCloser(strings.NewReader("system\n")), nil + }, + } + cfg := testOpsCleanupConfig() + cfg.Ops.Cleanup.MaxCatchupWindowsPerRun = 1 + svc := NewOpsCleanupService(repo, nil, db, nil, cfg, nil, &opsCleanupArchiveStub{ + create: func(context.Context, DataArchiveInput) (*BackupRecord, error) { + return &BackupRecord{Status: "completed"}, nil + }, + }) + loc, err := time.LoadLocation(cfg.Timezone) + if err != nil { + t.Fatal(err) + } + counts, err := svc.runCleanupOnceWithConfig(context.Background(), cfg.Ops.Cleanup, loc) + if err == nil || !strings.Contains(err.Error(), "error archive unavailable") { + t.Fatalf("error = %v, want error archive failure", err) + } + if counts.errorLogs != 0 || counts.systemLogs != 1 { + t.Fatalf("counts = %+v, want system logs to advance independently", counts) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatal(err) + } +} + +func TestOpsCleanupArchiveFailureDoesNotDelete(t *testing.T) { + db, mock, err := sqlmock.New() + if err != nil { + t.Fatal(err) + } + defer db.Close() + oldest := time.Date(2026, 6, 1, 12, 0, 0, 0, time.UTC) + mock.ExpectQuery(`SELECT MIN\(created_at\) FROM ops_error_logs`). + WithArgs(sqlmock.AnyArg()). + WillReturnRows(sqlmock.NewRows([]string{"min"}).AddRow(oldest)) + repo := &opsRepoMock{ExportErrorLogsFn: func(context.Context, *OpsErrorLogCleanupFilter) (io.ReadCloser, error) { + return io.NopCloser(strings.NewReader("row\n")), nil + }} + svc := NewOpsCleanupService(repo, nil, db, nil, testOpsCleanupConfig(), nil, &opsCleanupArchiveStub{ + create: func(context.Context, DataArchiveInput) (*BackupRecord, error) { + return nil, errors.New("upload failed") + }, + }) + + deleted, err := svc.cleanupOpsLogWindows(context.Background(), "ops_error_logs", "created_at", oldest.AddDate(0, 0, 10), svc.exportOpsErrorLogWindow) + if err == nil || !strings.Contains(err.Error(), "upload failed") { + t.Fatalf("error = %v, want upload failure", err) + } + if deleted != 0 { + t.Fatalf("deleted = %d, want 0", deleted) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatal(err) + } +} + +func TestOpsCleanupWindowsAdvanceByBoundedDay(t *testing.T) { + db, mock, err := sqlmock.New() + if err != nil { + t.Fatal(err) + } + defer db.Close() + first := time.Date(2026, 6, 1, 12, 0, 0, 0, time.UTC) + second := first.AddDate(0, 0, 1) + for _, oldest := range []time.Time{first, second} { + mock.ExpectQuery(`SELECT MIN\(created_at\) FROM ops_error_logs`). + WithArgs(sqlmock.AnyArg()). + WillReturnRows(sqlmock.NewRows([]string{"min"}).AddRow(oldest)) + mock.ExpectExec(`DELETE FROM ops_error_logs`). + WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg(), 5000). + WillReturnResult(sqlmock.NewResult(0, 1)) + } + var exported []opsCleanupWindow + repo := &opsRepoMock{ExportErrorLogsFn: func(_ context.Context, filter *OpsErrorLogCleanupFilter) (io.ReadCloser, error) { + exported = append(exported, opsCleanupWindow{Start: *filter.StartTime, End: *filter.EndTime}) + return io.NopCloser(strings.NewReader("row\n")), nil + }} + svc := NewOpsCleanupService(repo, nil, db, nil, testOpsCleanupConfig(), nil, &opsCleanupArchiveStub{ + create: func(context.Context, DataArchiveInput) (*BackupRecord, error) { + return &BackupRecord{Status: "completed"}, nil + }, + }) + + deleted, err := svc.cleanupOpsLogWindows(context.Background(), "ops_error_logs", "created_at", first.AddDate(0, 0, 10), svc.exportOpsErrorLogWindow) + if err != nil { + t.Fatal(err) + } + if deleted != 2 || len(exported) != 2 { + t.Fatalf("deleted=%d exported=%d, want 2/2", deleted, len(exported)) + } + if got := exported[0].End.Sub(exported[0].Start); got != 24*time.Hour { + t.Fatalf("first window duration = %s, want 24h", got) + } + if !exported[1].Start.Equal(exported[0].End) { + t.Fatalf("second window starts at %s, want %s", exported[1].Start, exported[0].End) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatal(err) + } +} + +func TestOpsCleanupArchiveTimeoutDoesNotDelete(t *testing.T) { + db, mock, err := sqlmock.New() + if err != nil { + t.Fatal(err) + } + defer db.Close() + oldest := time.Date(2026, 6, 1, 12, 0, 0, 0, time.UTC) + mock.ExpectQuery(`SELECT MIN\(created_at\) FROM ops_error_logs`). + WithArgs(sqlmock.AnyArg()). + WillReturnRows(sqlmock.NewRows([]string{"min"}).AddRow(oldest)) + repo := &opsRepoMock{ExportErrorLogsFn: func(context.Context, *OpsErrorLogCleanupFilter) (io.ReadCloser, error) { + return io.NopCloser(strings.NewReader("row\n")), nil + }} + svc := NewOpsCleanupService(repo, nil, db, nil, testOpsCleanupConfig(), nil, &opsCleanupArchiveStub{ + create: func(ctx context.Context, _ DataArchiveInput) (*BackupRecord, error) { + <-ctx.Done() + return nil, ctx.Err() + }, + }) + + deleted, err := svc.cleanupOpsLogWindows(context.Background(), "ops_error_logs", "created_at", oldest.AddDate(0, 0, 10), svc.exportOpsErrorLogWindow) + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("error = %v, want deadline exceeded", err) + } + if deleted != 0 { + t.Fatalf("deleted = %d, want 0", deleted) + } +} + +func TestOpsCleanupZeroRetentionDisablesTargets(t *testing.T) { + db, mock, err := sqlmock.New() + if err != nil { + t.Fatal(err) + } + defer db.Close() + cfg := testOpsCleanupConfig() + cfg.Ops.Cleanup.ErrorLogRetentionDays = 0 + svc := NewOpsCleanupService(&opsRepoMock{}, nil, db, nil, cfg, nil, nil) + counts, err := svc.runCleanupOnce(context.Background()) + if err != nil { + t.Fatal(err) + } + if counts != (opsCleanupDeletedCounts{}) { + t.Fatalf("counts = %+v, want zero", counts) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatal(err) + } +} + +func TestOpsCleanupArchivesBothLogTablesBeforeAuxiliaryDeletes(t *testing.T) { + db, mock, err := sqlmock.New() + if err != nil { + t.Fatal(err) + } + defer db.Close() + + oldest := time.Now().UTC().AddDate(0, 0, -40) + mock.ExpectQuery(`SELECT MIN\(created_at\) FROM ops_error_logs`). + WithArgs(sqlmock.AnyArg()). + WillReturnRows(sqlmock.NewRows([]string{"min"}).AddRow(oldest)) + mock.ExpectExec(`DELETE FROM ops_error_logs`). + WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg(), 5000). + WillReturnResult(sqlmock.NewResult(0, 1)) + mock.ExpectQuery(`SELECT MIN\(created_at\) FROM ops_system_logs`). + WithArgs(sqlmock.AnyArg()). + WillReturnRows(sqlmock.NewRows([]string{"min"}).AddRow(oldest)) + mock.ExpectExec(`DELETE FROM ops_system_logs`). + WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg(), 5000). + WillReturnResult(sqlmock.NewResult(0, 1)) + for _, table := range []string{"ops_retry_attempts", "ops_alert_events", "ops_system_log_cleanup_audits"} { + mock.ExpectExec(`DELETE FROM `+table). + WithArgs(sqlmock.AnyArg(), 5000). + WillReturnResult(sqlmock.NewResult(0, 0)) + } + + repo := &opsRepoMock{ + ExportErrorLogsFn: func(context.Context, *OpsErrorLogCleanupFilter) (io.ReadCloser, error) { + return io.NopCloser(strings.NewReader("error\n")), nil + }, + ExportSystemLogsFn: func(context.Context, *OpsSystemLogCleanupFilter) (io.ReadCloser, error) { + return io.NopCloser(strings.NewReader("system\n")), nil + }, + } + cfg := testOpsCleanupConfig() + cfg.Ops.Cleanup.MaxCatchupWindowsPerRun = 1 + svc := NewOpsCleanupService(repo, nil, db, nil, cfg, nil, &opsCleanupArchiveStub{ + create: func(context.Context, DataArchiveInput) (*BackupRecord, error) { + return &BackupRecord{Status: "completed"}, nil + }, + }) + + counts, err := svc.runCleanupOnce(context.Background()) + if err != nil { + t.Fatal(err) + } + if counts.errorLogs != 1 || counts.systemLogs != 1 { + t.Fatalf("counts = %+v, want one archived/deleted row per log table", counts) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatal(err) + } +} diff --git a/backend/internal/service/ops_service.go b/backend/internal/service/ops_service.go index cd3974a00..cfcd35160 100644 --- a/backend/internal/service/ops_service.go +++ b/backend/internal/service/ops_service.go @@ -15,6 +15,10 @@ import ( var ErrOpsDisabled = infraerrors.NotFound("OPS_DISABLED", "Ops monitoring is disabled") +type opsDataRetentionSettingsApplier interface { + ReconcileDataRetentionSettings(ctx context.Context) error +} + const ( opsMaxStoredRequestBodyBytes = 256 * 1024 opsMaxStoredErrorBodyBytes = 20 * 1024 @@ -54,6 +58,13 @@ type OpsService struct { geminiCompatService *GeminiMessagesCompatService antigravityGatewayService *AntigravityGatewayService systemLogSink *OpsSystemLogSink + cleanupSettingsApplier opsDataRetentionSettingsApplier +} + +func (s *OpsService) setCleanupSettingsApplier(applier opsDataRetentionSettingsApplier) { + if s != nil { + s.cleanupSettingsApplier = applier + } } func NewOpsService( diff --git a/backend/internal/service/ops_settings.go b/backend/internal/service/ops_settings.go index ecc3a94b7..b3aba5c54 100644 --- a/backend/internal/service/ops_settings.go +++ b/backend/internal/service/ops_settings.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "errors" + "fmt" "strings" "time" ) @@ -359,8 +360,8 @@ func (s *OpsService) UpdateOpsAlertRuntimeSettings(ctx context.Context, cfg *Ops func defaultOpsAdvancedSettings() *OpsAdvancedSettings { return &OpsAdvancedSettings{ DataRetention: OpsDataRetentionSettings{ - CleanupEnabled: false, - CleanupSchedule: "0 2 * * *", + CleanupEnabled: true, + CleanupSchedule: "0 4 * * *", ErrorLogRetentionDays: 30, MinuteMetricsRetentionDays: 30, HourlyMetricsRetentionDays: 30, @@ -385,9 +386,9 @@ func normalizeOpsAdvancedSettings(cfg *OpsAdvancedSettings) { } cfg.DataRetention.CleanupSchedule = strings.TrimSpace(cfg.DataRetention.CleanupSchedule) if cfg.DataRetention.CleanupSchedule == "" { - cfg.DataRetention.CleanupSchedule = "0 2 * * *" + cfg.DataRetention.CleanupSchedule = "0 4 * * *" } - // 保留天数：0 表示每次定时清理全部（清空所有），> 0 表示按天数保留； + // 保留天数：0 表示禁用对应目标，> 0 表示按自然日保留； // 仅在拿到非法的负数时回填默认值，避免覆盖用户主动设的 0。 if cfg.DataRetention.ErrorLogRetentionDays < 0 { cfg.DataRetention.ErrorLogRetentionDays = 30 @@ -408,18 +409,35 @@ func validateOpsAdvancedSettings(cfg *OpsAdvancedSettings) error { if cfg == nil { return errors.New("invalid config") } - // 保留天数：0 表示每次清理全部，1-365 表示按天数保留。 - if cfg.DataRetention.ErrorLogRetentionDays < 0 || cfg.DataRetention.ErrorLogRetentionDays > 365 { + if err := validateOpsDataRetentionSettings(cfg.DataRetention); err != nil { + return err + } + if cfg.AutoRefreshIntervalSec < 15 || cfg.AutoRefreshIntervalSec > 300 { + return errors.New("auto_refresh_interval_seconds must be between 15 and 300") + } + return nil +} + +func validateOpsDataRetentionSettings(cfg OpsDataRetentionSettings) error { + // 保留天数：0 表示禁用对应目标，1-365 表示按自然日保留。 + if cfg.ErrorLogRetentionDays < 0 || cfg.ErrorLogRetentionDays > 365 { return errors.New("error_log_retention_days must be between 0 and 365") } - if cfg.DataRetention.MinuteMetricsRetentionDays < 0 || cfg.DataRetention.MinuteMetricsRetentionDays > 365 { + if cfg.MinuteMetricsRetentionDays < 0 || cfg.MinuteMetricsRetentionDays > 365 { return errors.New("minute_metrics_retention_days must be between 0 and 365") } - if cfg.DataRetention.HourlyMetricsRetentionDays < 0 || cfg.DataRetention.HourlyMetricsRetentionDays > 365 { + if cfg.HourlyMetricsRetentionDays < 0 || cfg.HourlyMetricsRetentionDays > 365 { return errors.New("hourly_metrics_retention_days must be between 0 and 365") } - if cfg.AutoRefreshIntervalSec < 15 || cfg.AutoRefreshIntervalSec > 300 { - return errors.New("auto_refresh_interval_seconds must be between 15 and 300") + if !cfg.CleanupEnabled { + return nil + } + schedule := strings.TrimSpace(cfg.CleanupSchedule) + if schedule == "" { + return errors.New("cleanup_schedule is required when cleanup is enabled") + } + if _, err := opsCleanupCronParser.Parse(schedule); err != nil { + return fmt.Errorf("invalid cleanup_schedule: %w", err) } return nil } @@ -446,7 +464,7 @@ func (s *OpsService) GetOpsAdvancedSettings(ctx context.Context) (*OpsAdvancedSe cfg := defaultOpsAdvancedSettings() if err := json.Unmarshal([]byte(raw), cfg); err != nil { - return defaultCfg, nil + return nil, fmt.Errorf("decode ops advanced settings: %w", err) } normalizeOpsAdvancedSettings(cfg) @@ -476,6 +494,14 @@ func (s *OpsService) UpdateOpsAdvancedSettings(ctx context.Context, cfg *OpsAdva if err := s.settingRepo.Set(ctx, SettingKeyOpsAdvancedSettings, string(raw)); err != nil { return nil, err } + if s.cleanupSettingsApplier != nil { + reconcileCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + err := s.cleanupSettingsApplier.ReconcileDataRetentionSettings(reconcileCtx) + cancel() + if err != nil { + return nil, fmt.Errorf("apply ops cleanup settings: %w", err) + } + } updated := &OpsAdvancedSettings{} _ = json.Unmarshal(raw, updated) diff --git a/backend/internal/service/ops_settings_advanced_test.go b/backend/internal/service/ops_settings_advanced_test.go index 06cc545bb..f65128c80 100644 --- a/backend/internal/service/ops_settings_advanced_test.go +++ b/backend/internal/service/ops_settings_advanced_test.go @@ -20,11 +20,84 @@ func TestGetOpsAdvancedSettings_DefaultHidesOpenAITokenStats(t *testing.T) { if !cfg.DisplayAlertEvents { t.Fatalf("DisplayAlertEvents = false, want true by default") } + if !cfg.DataRetention.CleanupEnabled || cfg.DataRetention.CleanupSchedule != "0 4 * * *" { + t.Fatalf("cleanup default = %v/%q, want enabled at 04:00", cfg.DataRetention.CleanupEnabled, cfg.DataRetention.CleanupSchedule) + } if repo.setCalls != 1 { t.Fatalf("expected defaults to be persisted once, got %d", repo.setCalls) } } +type opsDataRetentionApplierStub struct { + called bool + err error +} + +func (s *opsDataRetentionApplierStub) ReconcileDataRetentionSettings(context.Context) error { + s.called = true + return s.err +} + +func TestUpdateOpsAdvancedSettingsAppliesCleanupScheduleImmediately(t *testing.T) { + repo := newRuntimeSettingRepoStub() + applier := &opsDataRetentionApplierStub{} + svc := &OpsService{settingRepo: repo, cleanupSettingsApplier: applier} + cfg := defaultOpsAdvancedSettings() + cfg.DataRetention.CleanupSchedule = "15 4 * * *" + cfg.DataRetention.ErrorLogRetentionDays = 0 + + if _, err := svc.UpdateOpsAdvancedSettings(context.Background(), cfg); err != nil { + t.Fatal(err) + } + if !applier.called { + t.Fatal("cleanup settings were not reconciled") + } + if repo.setCalls != 1 { + t.Fatalf("settings persisted %d times, want 1", repo.setCalls) + } +} + +func TestUpdateOpsAdvancedSettingsAllowsDisableWithHiddenInvalidCron(t *testing.T) { + repo := newRuntimeSettingRepoStub() + applier := &opsDataRetentionApplierStub{} + svc := &OpsService{settingRepo: repo, cleanupSettingsApplier: applier} + cfg := defaultOpsAdvancedSettings() + cfg.DataRetention.CleanupEnabled = false + cfg.DataRetention.CleanupSchedule = "invalid cron" + + if _, err := svc.UpdateOpsAdvancedSettings(context.Background(), cfg); err != nil { + t.Fatalf("disabled cleanup must be saveable even with a hidden stale cron: %v", err) + } + if repo.setCalls != 1 || !applier.called { + t.Fatalf("disabled settings not persisted/reconciled: setCalls=%d called=%v", repo.setCalls, applier.called) + } +} + +func TestGetOpsAdvancedSettingsRejectsCorruptedJSON(t *testing.T) { + repo := newRuntimeSettingRepoStub() + repo.values[SettingKeyOpsAdvancedSettings] = "{invalid" + svc := &OpsService{settingRepo: repo} + + if _, err := svc.GetOpsAdvancedSettings(context.Background()); err == nil { + t.Fatal("corrupted advanced settings must not be presented as healthy defaults") + } +} + +func TestUpdateOpsAdvancedSettingsRejectsInvalidCleanupCronBeforePersist(t *testing.T) { + repo := newRuntimeSettingRepoStub() + applier := &opsDataRetentionApplierStub{} + svc := &OpsService{settingRepo: repo, cleanupSettingsApplier: applier} + cfg := defaultOpsAdvancedSettings() + cfg.DataRetention.CleanupSchedule = "invalid cron" + + if _, err := svc.UpdateOpsAdvancedSettings(context.Background(), cfg); err == nil { + t.Fatal("invalid cleanup cron must be rejected") + } + if repo.setCalls != 0 || applier.called { + t.Fatalf("invalid settings must not persist or apply: setCalls=%d called=%v", repo.setCalls, applier.called) + } +} + func TestUpdateOpsAdvancedSettings_PersistsOpenAITokenStatsVisibility(t *testing.T) { repo := newRuntimeSettingRepoStub() svc := &OpsService{settingRepo: repo} diff --git a/backend/internal/service/scheduler_snapshot_hydration_test.go b/backend/internal/service/scheduler_snapshot_hydration_test.go index 478e89935..c8e191beb 100644 --- a/backend/internal/service/scheduler_snapshot_hydration_test.go +++ b/backend/internal/service/scheduler_snapshot_hydration_test.go @@ -4,13 +4,18 @@ package service import ( "context" + "errors" "testing" "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" + "github.com/stretchr/testify/require" ) type snapshotHydrationCache struct { - snapshot []*Account - accounts map[int64]*Account + snapshot []*Account + accounts map[int64]*Account + accountErr error } func (c *snapshotHydrationCache) GetSnapshot(ctx context.Context, bucket SchedulerBucket) ([]*Account, bool, error) { @@ -22,6 +27,9 @@ func (c *snapshotHydrationCache) SetSnapshot(ctx context.Context, bucket Schedul } func (c *snapshotHydrationCache) GetAccount(ctx context.Context, accountID int64) (*Account, error) { + if c.accountErr != nil { + return nil, c.accountErr + } if c.accounts == nil { return nil, nil } @@ -163,3 +171,80 @@ func TestGatewaySelectAccountWithLoadAwareness_HydratesSelectedAccountFromSchedu t.Fatalf("expected hydrated api key, got %q", got) } } + +func TestSelectionResultHydrationOwnsAcquiredSlotLifecycle(t *testing.T) { + ownerUserID := int64(501) + requestCtx := context.WithValue(context.Background(), ctxkey.AuthenticatedUserID, ownerUserID+1) + + tests := []struct { + name string + account *Account + accountErr error + wantErr bool + }{ + { + name: "missing hydrated account", + wantErr: true, + }, + { + name: "hydration read error", + accountErr: errors.New("snapshot read failed"), + wantErr: true, + }, + { + name: "owned pending account is invisible to another user", + account: &Account{ + ID: 9, + OwnerUserID: &ownerUserID, + ShareMode: AccountShareModePublic, + ShareStatus: AccountShareStatusPending, + }, + wantErr: true, + }, + { + name: "visible account transfers release ownership", + account: &Account{ID: 9}, + }, + } + + for _, serviceName := range []string{"gateway", "openai"} { + serviceName := serviceName + for _, tt := range tests { + tt := tt + t.Run(serviceName+"/"+tt.name, func(t *testing.T) { + cache := &snapshotHydrationCache{ + accounts: map[int64]*Account{9: tt.account}, + accountErr: tt.accountErr, + } + snapshot := NewSchedulerSnapshotService(cache, nil, nil, nil, nil) + releaseCount := 0 + release := func() { releaseCount++ } + metadata := &Account{ID: 9} + + var ( + result *AccountSelectionResult + err error + ) + if serviceName == "gateway" { + result, err = (&GatewayService{schedulerSnapshot: snapshot}).newSelectionResult(requestCtx, metadata, true, release, nil) + } else { + result, err = (&OpenAIGatewayService{schedulerSnapshot: snapshot}).newSelectionResult(requestCtx, metadata, true, release, nil) + } + + if tt.wantErr { + require.Error(t, err) + require.Nil(t, result) + require.Equal(t, 1, releaseCount, "hydration failure must release the acquired slot exactly once") + return + } + + require.NoError(t, err) + require.NotNil(t, result) + require.Zero(t, releaseCount, "a successful result transfers release ownership to the caller") + require.NotNil(t, result.ReleaseFunc) + result.ReleaseFunc() + require.Equal(t, 1, releaseCount) + }) + } + } +} diff --git a/backend/internal/service/scheduler_snapshot_service.go b/backend/internal/service/scheduler_snapshot_service.go index ee30ddb56..beaaa3ab7 100644 --- a/backend/internal/service/scheduler_snapshot_service.go +++ b/backend/internal/service/scheduler_snapshot_service.go @@ -411,6 +411,9 @@ func (s *SchedulerSnapshotService) GetAccount(ctx context.Context, accountID int return account, nil } } + if s.accountRepo == nil { + return nil, nil + } if err := s.guardFallback(ctx); err != nil { return nil, err diff --git a/backend/internal/service/setting_service.go b/backend/internal/service/setting_service.go index 2b3195c95..564de0853 100644 --- a/backend/internal/service/setting_service.go +++ b/backend/internal/service/setting_service.go @@ -558,6 +558,7 @@ func (s *SettingService) GetAllSettings(ctx context.Context) (*SystemSettings, e result := s.parseSettings(settings) result.WithdrawalRateLimitWindowDays = withdrawalRateLimit.WindowDays result.WithdrawalRateLimitMax = withdrawalRateLimit.MaxRequests + result.WithdrawalRateLimitExemptAmount = withdrawalRateLimit.ExemptAmount accountLevels, err := parseOpenAIAccountLevelConfigsSetting(settings[SettingKeyOpenAIAccountLevels]) if err != nil { return nil, fmt.Errorf("parse openai account levels: %w", err) @@ -721,6 +722,7 @@ func (s *SettingService) GetPublicSettings(ctx context.Context) (*PublicSettings SettingKeyWithdrawalManagementEnabled, SettingKeyWithdrawalRateLimitWindowDays, SettingKeyWithdrawalRateLimitMax, + SettingKeyWithdrawalRateLimitExemptAmount, } settings, err := s.settingRepo.GetMultiple(ctx, keys) @@ -839,10 +841,11 @@ func (s *SettingService) GetPublicSettings(ctx context.Context) (*PublicSettings RiskControlEnabled: settings[SettingKeyRiskControlEnabled] == "true", - InvoiceManagementEnabled: settings[SettingKeyInvoiceManagementEnabled] == "true", - WithdrawalManagementEnabled: !isFalseSettingValue(settings[SettingKeyWithdrawalManagementEnabled]), - WithdrawalRateLimitWindowDays: withdrawalRateLimit.WindowDays, - WithdrawalRateLimitMax: withdrawalRateLimit.MaxRequests, + InvoiceManagementEnabled: settings[SettingKeyInvoiceManagementEnabled] == "true", + WithdrawalManagementEnabled: !isFalseSettingValue(settings[SettingKeyWithdrawalManagementEnabled]), + WithdrawalRateLimitWindowDays: withdrawalRateLimit.WindowDays, + WithdrawalRateLimitMax: withdrawalRateLimit.MaxRequests, + WithdrawalRateLimitExemptAmount: withdrawalRateLimit.ExemptAmount, }, nil } @@ -1023,6 +1026,7 @@ type PublicSettingsInjectionPayload struct { WithdrawalManagementEnabled bool `json:"withdrawal_management_enabled"` WithdrawalRateLimitWindowDays int `json:"withdrawal_rate_limit_window_days"` WithdrawalRateLimitMax int `json:"withdrawal_rate_limit_max"` + WithdrawalRateLimitExemptAmount float64 `json:"withdrawal_rate_limit_exempt_amount"` } // GetPublicSettingsForInjection returns public settings in a format suitable for HTML injection. @@ -1090,6 +1094,7 @@ func (s *SettingService) GetPublicSettingsForInjection(ctx context.Context) (any WithdrawalManagementEnabled: settings.WithdrawalManagementEnabled, WithdrawalRateLimitWindowDays: settings.WithdrawalRateLimitWindowDays, WithdrawalRateLimitMax: settings.WithdrawalRateLimitMax, + WithdrawalRateLimitExemptAmount: settings.WithdrawalRateLimitExemptAmount, }, nil } @@ -1847,18 +1852,24 @@ func (s *SettingService) buildSystemSettingsUpdates(ctx context.Context, setting updates[SettingKeyInvoiceManagementEnabled] = strconv.FormatBool(settings.InvoiceManagementEnabled) updates[SettingKeyWithdrawalManagementEnabled] = strconv.FormatBool(settings.WithdrawalManagementEnabled) withdrawalRateLimit := WithdrawalRateLimitConfig{ - WindowDays: settings.WithdrawalRateLimitWindowDays, - MaxRequests: settings.WithdrawalRateLimitMax, + WindowDays: settings.WithdrawalRateLimitWindowDays, + MaxRequests: settings.WithdrawalRateLimitMax, + ExemptAmount: settings.WithdrawalRateLimitExemptAmount, } - if withdrawalRateLimit.WindowDays == 0 && withdrawalRateLimit.MaxRequests == 0 { + if withdrawalRateLimit.WindowDays == 0 && withdrawalRateLimit.MaxRequests == 0 && withdrawalRateLimit.ExemptAmount == 0 { withdrawalRateLimit.WindowDays = WithdrawalRateLimitWindowDaysDefault + withdrawalRateLimit.ExemptAmount = WithdrawalRateLimitExemptAmountDefault settings.WithdrawalRateLimitWindowDays = withdrawalRateLimit.WindowDays + settings.WithdrawalRateLimitExemptAmount = withdrawalRateLimit.ExemptAmount } if err := ValidateWithdrawalRateLimitConfig(withdrawalRateLimit); err != nil { return nil, err } + withdrawalRateLimit.ExemptAmount, _ = normalizeWithdrawalAmount(withdrawalRateLimit.ExemptAmount) + settings.WithdrawalRateLimitExemptAmount = withdrawalRateLimit.ExemptAmount updates[SettingKeyWithdrawalRateLimitWindowDays] = strconv.Itoa(withdrawalRateLimit.WindowDays) updates[SettingKeyWithdrawalRateLimitMax] = strconv.Itoa(withdrawalRateLimit.MaxRequests) + updates[SettingKeyWithdrawalRateLimitExemptAmount] = strconv.FormatFloat(withdrawalRateLimit.ExemptAmount, 'f', 2, 64) // Claude Code version check updates[SettingKeyMinClaudeCodeVersion] = settings.MinClaudeCodeVersion @@ -2367,6 +2378,7 @@ func (s *SettingService) GetWithdrawalRateLimitConfig(ctx context.Context) (With settings, err := s.settingRepo.GetMultiple(ctx, []string{ SettingKeyWithdrawalRateLimitWindowDays, SettingKeyWithdrawalRateLimitMax, + SettingKeyWithdrawalRateLimitExemptAmount, }) if err != nil { return WithdrawalRateLimitConfig{}, fmt.Errorf("get withdrawal rate limit settings: %w", err) @@ -2376,8 +2388,9 @@ func (s *SettingService) GetWithdrawalRateLimitConfig(ctx context.Context) (With func parseWithdrawalRateLimitConfig(settings map[string]string) (WithdrawalRateLimitConfig, error) { config := WithdrawalRateLimitConfig{ - WindowDays: WithdrawalRateLimitWindowDaysDefault, - MaxRequests: WithdrawalRateLimitMaxDefault, + WindowDays: WithdrawalRateLimitWindowDaysDefault, + MaxRequests: WithdrawalRateLimitMaxDefault, + ExemptAmount: WithdrawalRateLimitExemptAmountDefault, } if raw := strings.TrimSpace(settings[SettingKeyWithdrawalRateLimitWindowDays]); raw != "" { value, err := strconv.Atoi(raw) @@ -2399,9 +2412,20 @@ func parseWithdrawalRateLimitConfig(settings map[string]string) (WithdrawalRateL } config.MaxRequests = value } + if raw := strings.TrimSpace(settings[SettingKeyWithdrawalRateLimitExemptAmount]); raw != "" { + value, err := strconv.ParseFloat(raw, 64) + if err != nil { + return WithdrawalRateLimitConfig{}, infraerrors.BadRequest( + "WITHDRAWAL_RATE_LIMIT_CONFIG_INVALID", + "withdrawal rate limit exempt amount must be a number", + ) + } + config.ExemptAmount = value + } if err := ValidateWithdrawalRateLimitConfig(config); err != nil { return WithdrawalRateLimitConfig{}, err } + config.ExemptAmount, _ = normalizeWithdrawalAmount(config.ExemptAmount) return config, nil } @@ -2851,10 +2875,11 @@ func (s *SettingService) InitializeDefaultSettings(ctx context.Context) error { SettingKeyAffiliateEnabled: "false", // Functional modules - SettingKeyInvoiceManagementEnabled: "false", - SettingKeyWithdrawalManagementEnabled: "true", - SettingKeyWithdrawalRateLimitWindowDays: strconv.Itoa(WithdrawalRateLimitWindowDaysDefault), - SettingKeyWithdrawalRateLimitMax: strconv.Itoa(WithdrawalRateLimitMaxDefault), + SettingKeyInvoiceManagementEnabled: "false", + SettingKeyWithdrawalManagementEnabled: "true", + SettingKeyWithdrawalRateLimitWindowDays: strconv.Itoa(WithdrawalRateLimitWindowDaysDefault), + SettingKeyWithdrawalRateLimitMax: strconv.Itoa(WithdrawalRateLimitMaxDefault), + SettingKeyWithdrawalRateLimitExemptAmount: strconv.FormatFloat(WithdrawalRateLimitExemptAmountDefault, 'f', 2, 64), // Claude Code version check (default: empty = disabled) SettingKeyMinClaudeCodeVersion: "", @@ -3275,6 +3300,7 @@ func (s *SettingService) parseSettings(settings map[string]string) *SystemSettin withdrawalRateLimit, _ := parseWithdrawalRateLimitConfig(settings) result.WithdrawalRateLimitWindowDays = withdrawalRateLimit.WindowDays result.WithdrawalRateLimitMax = withdrawalRateLimit.MaxRequests + result.WithdrawalRateLimitExemptAmount = withdrawalRateLimit.ExemptAmount result.RiskControlEnabled = settings[SettingKeyRiskControlEnabled] == "true" result.CyberSessionBlockEnabled = settings[SettingKeyCyberSessionBlockEnabled] == "true" if v, err := strconv.Atoi(strings.TrimSpace(settings[SettingKeyCyberSessionBlockTTLSeconds])); err == nil && v > 0 { diff --git a/backend/internal/service/settings_view.go b/backend/internal/service/settings_view.go index b705d7b5b..c17c0ad24 100644 --- a/backend/internal/service/settings_view.go +++ b/backend/internal/service/settings_view.go @@ -139,6 +139,7 @@ type SystemSettings struct { WithdrawalManagementEnabled bool WithdrawalRateLimitWindowDays int WithdrawalRateLimitMax int + WithdrawalRateLimitExemptAmount float64 CyberSessionBlockEnabled bool CyberSessionBlockTTLSeconds int AccountShareCommentReviewEnabled bool @@ -312,10 +313,11 @@ type PublicSettings struct { AffiliateEnabled bool `json:"affiliate_enabled"` // Functional module switches - InvoiceManagementEnabled bool `json:"invoice_management_enabled"` - WithdrawalManagementEnabled bool `json:"withdrawal_management_enabled"` - WithdrawalRateLimitWindowDays int `json:"withdrawal_rate_limit_window_days"` - WithdrawalRateLimitMax int `json:"withdrawal_rate_limit_max"` + InvoiceManagementEnabled bool `json:"invoice_management_enabled"` + WithdrawalManagementEnabled bool `json:"withdrawal_management_enabled"` + WithdrawalRateLimitWindowDays int `json:"withdrawal_rate_limit_window_days"` + WithdrawalRateLimitMax int `json:"withdrawal_rate_limit_max"` + WithdrawalRateLimitExemptAmount float64 `json:"withdrawal_rate_limit_exempt_amount"` // 风控中心功能开关 RiskControlEnabled bool `json:"risk_control_enabled"` diff --git a/backend/internal/service/token_refresh_service.go b/backend/internal/service/token_refresh_service.go index 93d3cffd4..714c633c1 100644 --- a/backend/internal/service/token_refresh_service.go +++ b/backend/internal/service/token_refresh_service.go @@ -38,9 +38,10 @@ type TokenRefreshService struct { privacyClientFactory PrivacyClientFactory proxyRepo ProxyRepository - stopCh chan struct{} - stopOnce sync.Once - wg sync.WaitGroup + runCtx context.Context + cancelRun context.CancelFunc + stopOnce sync.Once + wg sync.WaitGroup } // NewTokenRefreshService 创建token刷新服务 @@ -56,6 +57,7 @@ func NewTokenRefreshService( tempUnschedCache TempUnschedCache, grokOAuthServices ...*GrokOAuthService, ) *TokenRefreshService { + runCtx, cancelRun := context.WithCancel(context.Background()) s := &TokenRefreshService{ accountRepo: accountRepo, refreshPolicy: DefaultBackgroundRefreshPolicy(), @@ -63,7 +65,8 @@ func NewTokenRefreshService( cacheInvalidator: cacheInvalidator, schedulerCache: schedulerCache, tempUnschedCache: tempUnschedCache, - stopCh: make(chan struct{}), + runCtx: runCtx, + cancelRun: cancelRun, } openAIRefresher := NewOpenAITokenRefresher(openaiOAuthService, accountRepo) @@ -136,7 +139,7 @@ func (s *TokenRefreshService) Start() { // Stop 停止刷新服务（可安全多次调用） func (s *TokenRefreshService) Stop() { s.stopOnce.Do(func() { - close(s.stopCh) + s.cancelRun() }) s.wg.Wait() slog.Info("token_refresh.service_stopped") @@ -155,22 +158,28 @@ func (s *TokenRefreshService) refreshLoop() { ticker := time.NewTicker(checkInterval) defer ticker.Stop() - // 启动时立即执行一次检查 - s.processRefresh() + // 启动时立即执行一次检查。Stop 可能在 goroutine 真正运行前已被调用， + // 因此必须先检查上下文，避免关停期间反而启动一轮刷新。 + if s.runCtx.Err() != nil { + return + } + s.processRefresh(s.runCtx) for { select { case <-ticker.C: - s.processRefresh() - case <-s.stopCh: + s.processRefresh(s.runCtx) + case <-s.runCtx.Done(): return } } } // processRefresh 执行一次刷新检查 -func (s *TokenRefreshService) processRefresh() { - ctx := context.Background() +func (s *TokenRefreshService) processRefresh(ctx context.Context) { + if ctx.Err() != nil { + return + } // 计算刷新窗口 refreshWindow := time.Duration(s.cfg.RefreshBeforeExpiryHours * float64(time.Hour)) @@ -178,6 +187,9 @@ func (s *TokenRefreshService) processRefresh() { // 获取所有active状态的账号 accounts, err := s.listActiveAccounts(ctx, refreshWindow) if err != nil { + if ctx.Err() != nil { + return + } slog.Error("token_refresh.list_accounts_failed", "error", err) return } @@ -188,10 +200,16 @@ func (s *TokenRefreshService) processRefresh() { refreshed, failed, skipped := 0, 0, 0 for i := range accounts { + if ctx.Err() != nil { + return + } account := &accounts[i] // 遍历所有刷新器，找到能处理此账号的 for idx, refresher := range s.refreshers { + if ctx.Err() != nil { + return + } if !refresher.CanRefresh(account) { continue } @@ -213,6 +231,9 @@ func (s *TokenRefreshService) processRefresh() { // 执行刷新 if err := s.refreshWithRetry(ctx, account, refresher, executor, refreshWindow); err != nil { + if ctx.Err() != nil { + return + } if errors.Is(err, errRefreshSkipped) { skipped++ } else { @@ -223,7 +244,7 @@ func (s *TokenRefreshService) processRefresh() { ) failed++ } - } else { + } else if ctx.Err() == nil { slog.Info("token_refresh.account_refreshed", "account_id", account.ID, "account_name", account.Name, @@ -266,6 +287,10 @@ func (s *TokenRefreshService) refreshWithRetry(ctx context.Context, account *Acc var lastErr error for attempt := 1; attempt <= s.cfg.MaxRetries; attempt++ { + if err := ctx.Err(); err != nil { + return err + } + var newCredentials map[string]any var err error @@ -295,8 +320,17 @@ func (s *TokenRefreshService) refreshWithRetry(ctx context.Context, account *Acc } } + // 关停取消是进程生命周期事件，不是账号刷新失败。 + // 必须在错误分类和状态写入前返回，避免停机将账号误标为 error/临时不可调度。 + if ctxErr := ctx.Err(); ctxErr != nil { + return ctxErr + } + if err == nil { s.postRefreshActions(ctx, account) + if ctxErr := ctx.Err(); ctxErr != nil { + return ctxErr + } return nil } @@ -324,9 +358,23 @@ func (s *TokenRefreshService) refreshWithRetry(ctx context.Context, account *Acc if attempt < s.cfg.MaxRetries { // 指数退避：2^(attempt-1) * baseSeconds backoff := time.Duration(s.cfg.RetryBackoffSeconds) * time.Second * time.Duration(1<<(attempt-1)) - time.Sleep(backoff) + timer := time.NewTimer(backoff) + select { + case <-timer.C: + case <-ctx.Done(): + if !timer.Stop() { + select { + case <-timer.C: + default: + } + } + return ctx.Err() + } } } + if err := ctx.Err(); err != nil { + return err + } // 可重试错误耗尽：临时标记账号不可调度，避免请求路径反复命中已知失败的账号 slog.Warn("token_refresh.retry_exhausted", @@ -356,6 +404,9 @@ func (s *TokenRefreshService) refreshWithRetry(ctx context.Context, account *Acc // postRefreshActions 刷新成功后的后续动作（清除错误状态、缓存失效、调度器同步等） func (s *TokenRefreshService) postRefreshActions(ctx context.Context, account *Account) { + if ctx.Err() != nil { + return + } // Antigravity 账户：如果之前是因为缺少 project_id 而标记为 error，现在成功获取到了，清除错误状态 if account.Platform == PlatformAntigravity && account.Status == StatusError && @@ -369,6 +420,9 @@ func (s *TokenRefreshService) postRefreshActions(ctx context.Context, account *A slog.Info("token_refresh.cleared_missing_project_id_error", "account_id", account.ID) } } + if ctx.Err() != nil { + return + } // 刷新成功后清除临时不可调度状态（处理 OAuth 401 恢复场景） if account.TempUnschedulableUntil != nil && time.Now().Before(*account.TempUnschedulableUntil) { if clearErr := s.accountRepo.ClearTempUnschedulable(ctx, account.ID); clearErr != nil { @@ -389,6 +443,9 @@ func (s *TokenRefreshService) postRefreshActions(ctx context.Context, account *A } } } + if ctx.Err() != nil { + return + } // 对所有 OAuth 账号调用缓存失效（InvalidateToken 内部根据平台判断是否需要处理） if s.cacheInvalidator != nil && account.Type == AccountTypeOAuth { if err := s.cacheInvalidator.InvalidateToken(ctx, account); err != nil { @@ -400,6 +457,9 @@ func (s *TokenRefreshService) postRefreshActions(ctx context.Context, account *A slog.Debug("token_refresh.token_cache_invalidated", "account_id", account.ID) } } + if ctx.Err() != nil { + return + } // 同步更新调度器缓存，确保调度获取的 Account 对象包含最新的 credentials if s.schedulerCache != nil { if err := s.schedulerCache.SetAccount(ctx, account); err != nil { @@ -411,8 +471,14 @@ func (s *TokenRefreshService) postRefreshActions(ctx context.Context, account *A slog.Debug("token_refresh.scheduler_cache_synced", "account_id", account.ID) } } + if ctx.Err() != nil { + return + } // OpenAI OAuth: 刷新成功后，检查是否已设置 privacy_mode，未设置则尝试关闭训练数据共享 s.ensureOpenAIPrivacy(ctx, account) + if ctx.Err() != nil { + return + } // Antigravity OAuth: 刷新成功后，检查是否已设置 privacy_mode，未设置则调用 setUserSettings s.ensureAntigravityPrivacy(ctx, account) } diff --git a/backend/internal/service/token_refresh_service_test.go b/backend/internal/service/token_refresh_service_test.go index 0d097d79b..e35bfb0d6 100644 --- a/backend/internal/service/token_refresh_service_test.go +++ b/backend/internal/service/token_refresh_service_test.go @@ -6,6 +6,7 @@ import ( "context" "errors" "reflect" + "sync" "testing" "time" @@ -141,6 +142,140 @@ func (r *tokenRefresherStub) CacheKey(account *Account) string { return "test:stub:" + account.Platform } +type blockingOAuthRefreshCandidateRepo struct { + tokenRefreshAccountRepo + listStarted chan struct{} + startOnce sync.Once +} + +func (r *blockingOAuthRefreshCandidateRepo) ListOAuthRefreshCandidates(ctx context.Context, _ time.Duration) ([]Account, error) { + r.startOnce.Do(func() { close(r.listStarted) }) + <-ctx.Done() + return nil, ctx.Err() +} + +type staticOAuthRefreshCandidateRepo struct { + tokenRefreshAccountRepo + candidates []Account +} + +func (r *staticOAuthRefreshCandidateRepo) ListOAuthRefreshCandidates(context.Context, time.Duration) ([]Account, error) { + return r.candidates, nil +} + +// observedDoneContext 让测试能精确知道刷新协程已进入退避 select， +// 避免依赖 sleep 猜测协程调度时序。 +type observedDoneContext struct { + context.Context + done chan struct{} + doneObserved chan struct{} + observeOnce sync.Once + cancelOnce sync.Once +} + +func newObservedDoneContext() *observedDoneContext { + return &observedDoneContext{ + Context: context.Background(), + done: make(chan struct{}), + doneObserved: make(chan struct{}), + } +} + +func (c *observedDoneContext) Done() <-chan struct{} { + c.observeOnce.Do(func() { close(c.doneObserved) }) + return c.done +} + +func (c *observedDoneContext) Err() error { + select { + case <-c.done: + return context.Canceled + default: + return nil + } +} + +func (c *observedDoneContext) cancel() { + c.cancelOnce.Do(func() { close(c.done) }) +} + +func requireTokenRefreshSignal(t *testing.T, signal <-chan struct{}, description string) { + t.Helper() + timer := time.NewTimer(time.Second) + defer timer.Stop() + select { + case <-signal: + case <-timer.C: + t.Fatalf("timed out waiting for %s", description) + } +} + +func requireTokenRefreshStop(t *testing.T, service *TokenRefreshService) { + t.Helper() + stopped := make(chan struct{}) + go func() { + service.Stop() + close(stopped) + }() + requireTokenRefreshSignal(t, stopped, "TokenRefreshService.Stop") +} + +func TestTokenRefreshService_StopCancelsBlockingCandidateList(t *testing.T) { + repo := &blockingOAuthRefreshCandidateRepo{listStarted: make(chan struct{})} + cfg := &config.Config{ + TokenRefresh: config.TokenRefreshConfig{ + Enabled: true, + CheckIntervalMinutes: 5, + MaxRetries: 3, + RetryBackoffSeconds: 2, + RefreshBeforeExpiryHours: 0.5, + }, + } + service := NewTokenRefreshService(repo, nil, nil, nil, nil, nil, nil, cfg, nil) + service.Start() + requireTokenRefreshSignal(t, repo.listStarted, "OAuth refresh candidate query to start") + + requireTokenRefreshStop(t, service) + require.Zero(t, repo.setErrorCalls) + require.Zero(t, repo.setTempUnschedCalls) +} + +func TestTokenRefreshService_StopInterruptsRetryBackoffWithoutAccountPenalty(t *testing.T) { + repo := &staticOAuthRefreshCandidateRepo{ + candidates: []Account{{ + ID: 1001, + Platform: PlatformGemini, + Type: AccountTypeOAuth, + }}, + } + cfg := &config.Config{ + TokenRefresh: config.TokenRefreshConfig{ + Enabled: true, + CheckIntervalMinutes: 5, + MaxRetries: 3, + RetryBackoffSeconds: 60, + RefreshBeforeExpiryHours: 0.5, + }, + } + service := NewTokenRefreshService(repo, nil, nil, nil, nil, nil, nil, cfg, nil) + refresher := &tokenRefresherStub{err: errors.New("transient token endpoint failure")} + service.refreshers = []TokenRefresher{refresher} + service.executors = []OAuthRefreshExecutor{refresher} + + // 用可观测 context 替换构造器的 context，只用于确认已真正进入 60s 退避等待。 + service.cancelRun() + runCtx := newObservedDoneContext() + service.runCtx = runCtx + service.cancelRun = runCtx.cancel + + service.Start() + requireTokenRefreshSignal(t, runCtx.doneObserved, "retry backoff wait to start") + requireTokenRefreshStop(t, service) + + require.Zero(t, repo.setErrorCalls, "shutdown cancellation must not mark the account as error") + require.Zero(t, repo.setTempUnschedCalls, "shutdown cancellation must not temporarily unschedule the account") +} + func TestTokenRefreshService_RefreshWithRetry_InvalidatesCache(t *testing.T) { repo := &tokenRefreshAccountRepo{} invalidator := &tokenCacheInvalidatorStub{} diff --git a/backend/internal/service/update_service.go b/backend/internal/service/update_service.go index 34ad46102..decbd1783 100644 --- a/backend/internal/service/update_service.go +++ b/backend/internal/service/update_service.go @@ -17,6 +17,8 @@ import ( "strconv" "strings" "time" + + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" ) const ( @@ -32,6 +34,11 @@ const ( maxDownloadSize = 500 * 1024 * 1024 ) +var ErrInPlaceUpdateDisabled = infraerrors.Forbidden( + "IN_PLACE_UPDATE_DISABLED", + "in-place update is disabled for this customized deployment; use the Pixel release and symlink deployment workflow", +) + // UpdateCache defines cache operations for update service type UpdateCache interface { GetUpdateInfo(ctx context.Context) (string, error) @@ -51,6 +58,7 @@ type UpdateService struct { githubClient GitHubReleaseClient currentVersion string buildType string // "source" for manual builds, "release" for CI builds + inPlaceAllowed bool } // NewUpdateService creates a new UpdateService @@ -60,6 +68,7 @@ func NewUpdateService(cache UpdateCache, githubClient GitHubReleaseClient, versi githubClient: githubClient, currentVersion: version, buildType: buildType, + inPlaceAllowed: envBool("PIXEL_ALLOW_UPSTREAM_IN_PLACE_UPDATE"), } } @@ -72,6 +81,7 @@ type UpdateInfo struct { Cached bool `json:"cached"` Warning string `json:"warning,omitempty"` BuildType string `json:"build_type"` // "source" or "release" + InPlaceAllowed bool `json:"in_place_update_allowed"` } // ReleaseInfo contains GitHub release details @@ -129,6 +139,7 @@ func (s *UpdateService) CheckUpdate(ctx context.Context, force bool) (*UpdateInf HasUpdate: false, Warning: err.Error(), BuildType: s.buildType, + InPlaceAllowed: s.inPlaceAllowed, }, nil } @@ -140,6 +151,9 @@ func (s *UpdateService) CheckUpdate(ctx context.Context, force bool) (*UpdateInf // PerformUpdate downloads and applies the update // Uses atomic file replacement pattern for safe in-place updates func (s *UpdateService) PerformUpdate(ctx context.Context) error { + if !s.inPlaceAllowed { + return ErrInPlaceUpdateDisabled + } info, err := s.CheckUpdate(ctx, true) if err != nil { return err @@ -251,6 +265,9 @@ func (s *UpdateService) PerformUpdate(ctx context.Context) error { // Rollback restores the previous version func (s *UpdateService) Rollback() error { + if !s.inPlaceAllowed { + return ErrInPlaceUpdateDisabled + } exePath, err := os.Executable() if err != nil { return fmt.Errorf("failed to get executable path: %w", err) @@ -301,8 +318,9 @@ func (s *UpdateService) fetchLatestRelease(ctx context.Context) (*UpdateInfo, er HTMLURL: release.HTMLURL, Assets: assets, }, - Cached: false, - BuildType: s.buildType, + Cached: false, + BuildType: s.buildType, + InPlaceAllowed: s.inPlaceAllowed, }, nil } @@ -493,9 +511,15 @@ func (s *UpdateService) getFromCache(ctx context.Context) (*UpdateInfo, error) { ReleaseInfo: cached.ReleaseInfo, Cached: true, BuildType: s.buildType, + InPlaceAllowed: s.inPlaceAllowed, }, nil } +func envBool(key string) bool { + value, err := strconv.ParseBool(strings.TrimSpace(os.Getenv(key))) + return err == nil && value +} + func (s *UpdateService) saveToCache(ctx context.Context, info *UpdateInfo) { cacheData := struct { Latest string `json:"latest"` diff --git a/backend/internal/service/update_service_security_test.go b/backend/internal/service/update_service_security_test.go new file mode 100644 index 000000000..d5f0aae52 --- /dev/null +++ b/backend/internal/service/update_service_security_test.go @@ -0,0 +1,15 @@ +package service + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestUpdateServiceCustomizedDeploymentDisablesInPlaceMutation(t *testing.T) { + svc := &UpdateService{inPlaceAllowed: false} + + require.ErrorIs(t, svc.PerformUpdate(context.Background()), ErrInPlaceUpdateDisabled) + require.ErrorIs(t, svc.Rollback(), ErrInPlaceUpdateDisabled) +} diff --git a/backend/internal/service/usage_log.go b/backend/internal/service/usage_log.go index f360f17d1..a74b8bd41 100644 --- a/backend/internal/service/usage_log.go +++ b/backend/internal/service/usage_log.go @@ -142,6 +142,8 @@ type UsageLog struct { CacheCreation5mTokens int `gorm:"column:cache_creation_5m_tokens"` CacheCreation1hTokens int `gorm:"column:cache_creation_1h_tokens"` + ImageInputTokens int + ImageInputCost float64 ImageOutputTokens int ImageOutputCost float64 diff --git a/backend/internal/service/usage_service.go b/backend/internal/service/usage_service.go index 0d26061b0..2f2016e2a 100644 --- a/backend/internal/service/usage_service.go +++ b/backend/internal/service/usage_service.go @@ -317,8 +317,8 @@ func (s *UsageService) GetAPIKeyDashboardStats(ctx context.Context, apiKeyID int } // GetUserUsageTrendByUserID returns per-user usage trend. -func (s *UsageService) GetUserUsageTrendByUserID(ctx context.Context, userID int64, startTime, endTime time.Time, granularity string) ([]usagestats.TrendDataPoint, error) { - trend, err := s.usageRepo.GetUserUsageTrendByUserID(ctx, userID, startTime, endTime, granularity) +func (s *UsageService) GetUserUsageTrendByUserID(ctx context.Context, userID int64, startTime, endTime time.Time, granularity string, location *time.Location) ([]usagestats.TrendDataPoint, error) { + trend, err := s.usageRepo.GetUserUsageTrendByUserID(ctx, userID, startTime, endTime, granularity, location) if err != nil { return nil, fmt.Errorf("get user usage trend: %w", err) } @@ -335,8 +335,8 @@ func (s *UsageService) GetUserModelStats(ctx context.Context, userID int64, star } // GetUserAccountSharingDashboard returns owned-account consumption and public-sharing settlement stats. -func (s *UsageService) GetUserAccountSharingDashboard(ctx context.Context, userID int64, startTime, endTime time.Time, granularity string) (*usagestats.AccountSharingDashboardStats, error) { - stats, err := s.usageRepo.GetUserAccountSharingDashboard(ctx, userID, startTime, endTime, granularity) +func (s *UsageService) GetUserAccountSharingDashboard(ctx context.Context, userID int64, startTime, endTime time.Time, granularity string, location *time.Location) (*usagestats.AccountSharingDashboardStats, error) { + stats, err := s.usageRepo.GetUserAccountSharingDashboard(ctx, userID, startTime, endTime, granularity, location) if err != nil { return nil, fmt.Errorf("get user account sharing dashboard: %w", err) } diff --git a/backend/internal/service/wire.go b/backend/internal/service/wire.go index 2f20d2973..e6a2788c3 100644 --- a/backend/internal/service/wire.go +++ b/backend/internal/service/wire.go @@ -341,14 +341,17 @@ func ProvideOpsAlertEvaluatorService( // channelMonitorSvc 让维护任务（聚合 + 历史/聚合软删）跟随 ops 清理 cron 一起跑， // 共享 leader lock + heartbeat。 func ProvideOpsCleanupService( + opsService *OpsService, opsRepo OpsRepository, + settingRepo SettingRepository, db *sql.DB, redisClient *redis.Client, cfg *config.Config, channelMonitorSvc *ChannelMonitorService, backupSvc *BackupService, ) *OpsCleanupService { - svc := NewOpsCleanupService(opsRepo, db, redisClient, cfg, channelMonitorSvc, backupSvc) + svc := NewOpsCleanupService(opsRepo, settingRepo, db, redisClient, cfg, channelMonitorSvc, backupSvc) + opsService.setCleanupSettingsApplier(svc) svc.Start() return svc } @@ -563,6 +566,7 @@ func ProvideAccountService( privateGroupProvisioner UserPrivateGroupProvisioner, systemNoticeService *SystemNoticeService, settingService *SettingService, + agentIdentityWSInvalidator *AgentIdentityWSInvalidatorProxy, ) *AccountService { svc := NewAccountService(accountRepo, groupRepo, userRepo, userSubRepo, proxyRepo) svc.SetAccountSharePolicyRepository(accountSharePolicyRepo) @@ -570,6 +574,7 @@ func ProvideAccountService( svc.SetUserPrivateGroupProvisioner(privateGroupProvisioner) svc.SetSystemNoticeService(systemNoticeService) svc.SetSettingService(settingService) + svc.SetAgentIdentityWSInvalidator(agentIdentityWSInvalidator) return svc } @@ -844,6 +849,7 @@ func ProvideAdminService( privacyClientFactory PrivacyClientFactory, privateGroupProvisioner UserPrivateGroupProvisioner, systemNoticeService *SystemNoticeService, + agentIdentityWSInvalidator *AgentIdentityWSInvalidatorProxy, ) AdminService { svc := NewAdminService( userRepo, @@ -866,7 +872,8 @@ func ProvideAdminService( privacyClientFactory, ) svc = SetAdminUserPrivateGroupProvisioner(svc, privateGroupProvisioner) - return SetAdminSystemNoticeService(svc, systemNoticeService) + svc = SetAdminSystemNoticeService(svc, systemNoticeService) + return SetAdminAgentIdentityWSInvalidator(svc, agentIdentityWSInvalidator) } // ProviderSet is the Wire provider set for all services diff --git a/backend/internal/service/withdrawal.go b/backend/internal/service/withdrawal.go index 49dce9b5b..4e4c23d3e 100644 --- a/backend/internal/service/withdrawal.go +++ b/backend/internal/service/withdrawal.go @@ -19,11 +19,12 @@ const ( WithdrawalMinimumAmount = 1.00 WithdrawalFirstFee = 0.10 - WithdrawalRateLimitWindowDaysDefault = 1 - WithdrawalRateLimitWindowDaysMin = 1 - WithdrawalRateLimitWindowDaysMax = 365 - WithdrawalRateLimitMaxDefault = 0 - WithdrawalRateLimitMaxAllowed = 1000 + WithdrawalRateLimitWindowDaysDefault = 1 + WithdrawalRateLimitWindowDaysMin = 1 + WithdrawalRateLimitWindowDaysMax = 365 + WithdrawalRateLimitMaxDefault = 0 + WithdrawalRateLimitMaxAllowed = 1000 + WithdrawalRateLimitExemptAmountDefault = 500.00 ) var ( @@ -40,8 +41,9 @@ var ( ) type WithdrawalRateLimitConfig struct { - WindowDays int - MaxRequests int + WindowDays int + MaxRequests int + ExemptAmount float64 } func ValidateWithdrawalRateLimitConfig(config WithdrawalRateLimitConfig) error { @@ -57,9 +59,23 @@ func ValidateWithdrawalRateLimitConfig(config WithdrawalRateLimitConfig) error { "withdrawal rate limit max must be between 0 and 1000", ) } + if _, ok := normalizeWithdrawalAmount(config.ExemptAmount); !ok || config.ExemptAmount < 0 { + return infraerrors.BadRequest( + "WITHDRAWAL_RATE_LIMIT_CONFIG_INVALID", + "withdrawal rate limit exempt amount must be non-negative and use at most two decimal places", + ) + } return nil } +// ExemptsAmount reports whether a withdrawal is outside the frequency limit. +// A zero threshold disables the exemption; an amount equal to the threshold is still limited. +func (config WithdrawalRateLimitConfig) ExemptsAmount(amount float64) bool { + threshold, thresholdOK := normalizeWithdrawalAmount(config.ExemptAmount) + normalizedAmount, amountOK := normalizeWithdrawalAmount(amount) + return thresholdOK && amountOK && threshold > 0 && normalizedAmount > threshold +} + func NewWithdrawalRateLimitExceededError(config WithdrawalRateLimitConfig) error { return infraerrors.TooManyRequests( "WITHDRAWAL_RATE_LIMIT_EXCEEDED", @@ -281,8 +297,9 @@ func (s *WithdrawalService) ensureEnabled(ctx context.Context) error { func (s *WithdrawalService) getRateLimitConfig(ctx context.Context) (WithdrawalRateLimitConfig, error) { if s == nil || s.settingService == nil { return WithdrawalRateLimitConfig{ - WindowDays: WithdrawalRateLimitWindowDaysDefault, - MaxRequests: WithdrawalRateLimitMaxDefault, + WindowDays: WithdrawalRateLimitWindowDaysDefault, + MaxRequests: WithdrawalRateLimitMaxDefault, + ExemptAmount: WithdrawalRateLimitExemptAmountDefault, }, nil } return s.settingService.GetWithdrawalRateLimitConfig(ctx) diff --git a/backend/internal/web/embed_on.go b/backend/internal/web/embed_on.go index 9b8b373dd..9de452ff6 100644 --- a/backend/internal/web/embed_on.go +++ b/backend/internal/web/embed_on.go @@ -7,9 +7,11 @@ import ( "context" "embed" "encoding/json" + htmlpkg "html" "io" "io/fs" "net/http" + "net/url" "os" "path/filepath" "strings" @@ -208,12 +210,62 @@ func (s *FrontendServer) injectSettings(settingsJSON []byte) []byte { headClose := []byte("") result := bytes.Replace(s.baseHTML, headClose, append(script, headClose...), 1) - // Replace with custom site name so the browser tab shows it immediately + // 在浏览器首次绘制前应用品牌名称和图标，避免静态默认品牌闪烁。 result = injectSiteTitle(result, settingsJSON) + result = injectSiteFavicon(result, settingsJSON) return result } +func injectSiteFavicon(html, settingsJSON []byte) []byte { + var cfg struct { + SiteLogo string `json:"site_logo"` + } + if err := json.Unmarshal(settingsJSON, &cfg); err != nil { + return html + } + + logoURL := safeBrandImageURL(cfg.SiteLogo) + if logoURL == "" { + return html + } + linkStart := bytes.Index(html, []byte(`<link rel="icon"`)) + if linkStart == -1 { + return html + } + linkEndOffset := bytes.IndexByte(html[linkStart:], '>') + if linkEndOffset == -1 { + return html + } + linkEnd := linkStart + linkEndOffset + 1 + replacement := []byte(`<link rel="icon" href="` + htmlpkg.EscapeString(logoURL) + `" />`) + + var buf bytes.Buffer + buf.Write(html[:linkStart]) + buf.Write(replacement) + buf.Write(html[linkEnd:]) + return buf.Bytes() +} + +func safeBrandImageURL(value string) string { + trimmed := strings.TrimSpace(value) + if trimmed == "" { + return "" + } + if strings.HasPrefix(trimmed, "/") && !strings.HasPrefix(trimmed, "//") { + return trimmed + } + if strings.HasPrefix(strings.ToLower(trimmed), "data:image/") { + return trimmed + } + parsed, err := url.Parse(trimmed) + if err != nil || parsed.User != nil || parsed.Host == "" || + (parsed.Scheme != "http" && parsed.Scheme != "https") { + return "" + } + return trimmed +} + // injectSiteTitle replaces the static <title> in HTML with the configured site name. // This ensures the browser tab shows the correct title before JS executes. func injectSiteTitle(html, settingsJSON []byte) []byte { @@ -231,7 +283,7 @@ func injectSiteTitle(html, settingsJSON []byte) []byte { return html } - newTitle := []byte("<title>" + cfg.SiteName + " - AI API Gateway") + newTitle := []byte("" + htmlpkg.EscapeString(cfg.SiteName) + " - AI API Gateway") var buf bytes.Buffer buf.Write(html[:titleStart]) buf.Write(newTitle) diff --git a/backend/internal/web/embed_test.go b/backend/internal/web/embed_test.go index 46771a899..f96fe1ab3 100644 --- a/backend/internal/web/embed_test.go +++ b/backend/internal/web/embed_test.go @@ -97,6 +97,24 @@ func TestInjectSiteTitle(t *testing.T) { assert.Contains(t, string(result), `
`) assert.Contains(t, string(result), "TestSite - AI API Gateway") }) + + t.Run("escapes_site_name", func(t *testing.T) { + html := []byte(`Pixel`) + result := injectSiteTitle(html, []byte(`{"site_name":""}`)) + + assert.NotContains(t, string(result), ``) + assert.Contains(t, string(result), `</title><script>alert(1)</script>`) + }) +} + +func TestInjectSiteFavicon(t *testing.T) { + html := []byte(``) + + result := injectSiteFavicon(html, []byte(`{"site_logo":"https://cdn.example.com/logo.png?a=1&b=2"}`)) + require.Contains(t, string(result), `href="https://cdn.example.com/logo.png?a=1&b=2"`) + + unsafeResult := injectSiteFavicon(html, []byte(`{"site_logo":"javascript:alert(1)"}`)) + require.Equal(t, string(html), string(unsafeResult)) } func TestReplaceNoncePlaceholder(t *testing.T) { diff --git a/backend/migrations/215_ops_daily_partition_shadow.sql b/backend/migrations/215_ops_daily_partition_shadow.sql new file mode 100644 index 000000000..1d52955e0 --- /dev/null +++ b/backend/migrations/215_ops_daily_partition_shadow.sql @@ -0,0 +1,364 @@ +-- OPS 日分区迁移的空影子基础设施。 +-- +-- 安全边界： +-- 1. 本迁移只创建空的分区父表和分区创建函数； +-- 2. 不创建任何日分区，不复制数据，不修改、附加或重命名正式表； +-- 3. 日分区必须由运维显式传入 UTC 零点后创建。 + +SET LOCAL lock_timeout = '2s'; + +DO $$ +DECLARE + relation_kind "char"; +BEGIN + SELECT c.relkind + INTO relation_kind + FROM pg_catalog.pg_class AS c + JOIN pg_catalog.pg_namespace AS n ON n.oid = c.relnamespace + WHERE n.nspname = 'public' + AND c.relname = 'ops_system_logs'; + + IF relation_kind IS DISTINCT FROM 'r'::"char" THEN + RAISE EXCEPTION 'public.ops_system_logs must exist as an ordinary table before creating its shadow parent'; + END IF; + + SELECT c.relkind + INTO relation_kind + FROM pg_catalog.pg_class AS c + JOIN pg_catalog.pg_namespace AS n ON n.oid = c.relnamespace + WHERE n.nspname = 'public' + AND c.relname = 'ops_error_logs'; + + IF relation_kind IS DISTINCT FROM 'r'::"char" THEN + RAISE EXCEPTION 'public.ops_error_logs must exist as an ordinary table before creating its shadow parent'; + END IF; +END $$; + +CREATE TABLE IF NOT EXISTS public.ops_system_logs_daily_shadow ( + LIKE public.ops_system_logs + INCLUDING DEFAULTS + INCLUDING GENERATED + INCLUDING IDENTITY + INCLUDING STORAGE + INCLUDING COMPRESSION + INCLUDING COMMENTS + INCLUDING CONSTRAINTS, + CONSTRAINT ops_system_logs_daily_shadow_pkey PRIMARY KEY (created_at, id) +) PARTITION BY RANGE (created_at); + +CREATE TABLE IF NOT EXISTS public.ops_error_logs_daily_shadow ( + LIKE public.ops_error_logs + INCLUDING DEFAULTS + INCLUDING GENERATED + INCLUDING IDENTITY + INCLUDING STORAGE + INCLUDING COMPRESSION + INCLUDING COMMENTS + INCLUDING CONSTRAINTS, + CONSTRAINT ops_error_logs_daily_shadow_pkey PRIMARY KEY (created_at, id) +) PARTITION BY RANGE (created_at); + +-- 双写和回填必须复用正式表已经分配的 id。影子表不允许自行推进正式序列。 +ALTER TABLE public.ops_system_logs_daily_shadow ALTER COLUMN id DROP DEFAULT; +ALTER TABLE public.ops_error_logs_daily_shadow ALTER COLUMN id DROP DEFAULT; + +DO $$ +DECLARE + parent_name text; + parent_oid oid; + source_name text; + source_oid oid; + partition_key text; + required_primary_key text; + id_index_name text; + id_attribute_number smallint; +BEGIN + FOREACH parent_name IN ARRAY ARRAY[ + 'ops_system_logs_daily_shadow', + 'ops_error_logs_daily_shadow' + ] + LOOP + source_name := pg_catalog.replace(parent_name, '_daily_shadow', ''); + + SELECT c.oid, pg_catalog.pg_get_partkeydef(c.oid) + INTO parent_oid, partition_key + FROM pg_catalog.pg_class AS c + JOIN pg_catalog.pg_namespace AS n ON n.oid = c.relnamespace + WHERE n.nspname = 'public' + AND c.relname = parent_name + AND c.relkind = 'p'; + + IF parent_oid IS NULL OR partition_key IS DISTINCT FROM 'RANGE (created_at)' THEN + RAISE EXCEPTION 'public.% must be a RANGE (created_at) partitioned table', parent_name; + END IF; + + SELECT c.oid + INTO source_oid + FROM pg_catalog.pg_class AS c + JOIN pg_catalog.pg_namespace AS n ON n.oid = c.relnamespace + WHERE n.nspname = 'public' + AND c.relname = source_name + AND c.relkind = 'r'; + + IF source_oid IS NULL THEN + RAISE EXCEPTION 'required source table public.% does not exist', source_name; + END IF; + + IF EXISTS ( + ( + SELECT a.attnum, + a.attname, + a.atttypid, + a.atttypmod, + a.attcollation, + a.attnotnull, + a.attidentity, + a.attgenerated, + CASE + WHEN a.attname = 'id' THEN '' + ELSE COALESCE(pg_catalog.pg_get_expr(d.adbin, d.adrelid), '') + END + FROM pg_catalog.pg_attribute AS a + LEFT JOIN pg_catalog.pg_attrdef AS d + ON d.adrelid = a.attrelid + AND d.adnum = a.attnum + WHERE a.attrelid = source_oid + AND a.attnum > 0 + AND NOT a.attisdropped + EXCEPT + SELECT a.attnum, + a.attname, + a.atttypid, + a.atttypmod, + a.attcollation, + a.attnotnull, + a.attidentity, + a.attgenerated, + CASE + WHEN a.attname = 'id' THEN '' + ELSE COALESCE(pg_catalog.pg_get_expr(d.adbin, d.adrelid), '') + END + FROM pg_catalog.pg_attribute AS a + LEFT JOIN pg_catalog.pg_attrdef AS d + ON d.adrelid = a.attrelid + AND d.adnum = a.attnum + WHERE a.attrelid = parent_oid + AND a.attnum > 0 + AND NOT a.attisdropped + ) + UNION ALL + ( + SELECT a.attnum, + a.attname, + a.atttypid, + a.atttypmod, + a.attcollation, + a.attnotnull, + a.attidentity, + a.attgenerated, + CASE + WHEN a.attname = 'id' THEN '' + ELSE COALESCE(pg_catalog.pg_get_expr(d.adbin, d.adrelid), '') + END + FROM pg_catalog.pg_attribute AS a + LEFT JOIN pg_catalog.pg_attrdef AS d + ON d.adrelid = a.attrelid + AND d.adnum = a.attnum + WHERE a.attrelid = parent_oid + AND a.attnum > 0 + AND NOT a.attisdropped + EXCEPT + SELECT a.attnum, + a.attname, + a.atttypid, + a.atttypmod, + a.attcollation, + a.attnotnull, + a.attidentity, + a.attgenerated, + CASE + WHEN a.attname = 'id' THEN '' + ELSE COALESCE(pg_catalog.pg_get_expr(d.adbin, d.adrelid), '') + END + FROM pg_catalog.pg_attribute AS a + LEFT JOIN pg_catalog.pg_attrdef AS d + ON d.adrelid = a.attrelid + AND d.adnum = a.attnum + WHERE a.attrelid = source_oid + AND a.attnum > 0 + AND NOT a.attisdropped + ) + ) THEN + RAISE EXCEPTION 'public.% columns/defaults do not match public.%', parent_name, source_name; + END IF; + + required_primary_key := parent_name || '_pkey'; + IF NOT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_constraint AS con + WHERE con.conrelid = parent_oid + AND con.contype = 'p' + AND con.conname = required_primary_key + AND pg_catalog.pg_get_constraintdef(con.oid, true) = 'PRIMARY KEY (created_at, id)' + ) THEN + RAISE EXCEPTION 'public.% must have PRIMARY KEY (created_at, id)', parent_name; + END IF; + + id_index_name := parent_name || '_id_idx'; + IF pg_catalog.to_regclass('public.' || id_index_name) IS NULL THEN + IF EXISTS ( + SELECT 1 + FROM pg_catalog.pg_inherits + WHERE inhparent = parent_oid + ) THEN + RAISE EXCEPTION 'refusing to build missing index % after partitions exist', id_index_name; + END IF; + + EXECUTE pg_catalog.format( + 'CREATE INDEX %I ON public.%I (id)', + id_index_name, + parent_name + ); + END IF; + + SELECT a.attnum + INTO id_attribute_number + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = parent_oid + AND a.attname = 'id' + AND NOT a.attisdropped; + + IF NOT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_class AS idx + JOIN pg_catalog.pg_namespace AS idx_ns ON idx_ns.oid = idx.relnamespace + JOIN pg_catalog.pg_index AS ind ON ind.indexrelid = idx.oid + JOIN pg_catalog.pg_am AS am ON am.oid = idx.relam + WHERE idx_ns.nspname = 'public' + AND idx.relname = id_index_name + AND ind.indrelid = parent_oid + AND ind.indisvalid + AND ind.indisready + AND NOT ind.indisunique + AND ind.indnkeyatts = 1 + AND ind.indkey[0] = id_attribute_number + AND ind.indexprs IS NULL + AND ind.indpred IS NULL + AND am.amname = 'btree' + ) THEN + RAISE EXCEPTION 'public.% must be a valid non-unique btree index on public.%(id)', id_index_name, parent_name; + END IF; + END LOOP; +END $$; + +COMMENT ON TABLE public.ops_system_logs_daily_shadow IS + 'Empty UTC daily-partitioned shadow for an operator-controlled ops_system_logs migration; writes must provide the source id; never auto-backfilled or auto-switched.'; +COMMENT ON TABLE public.ops_error_logs_daily_shadow IS + 'Empty UTC daily-partitioned shadow for an operator-controlled ops_error_logs migration; writes must provide the source id; never auto-backfilled or auto-switched.'; + +CREATE OR REPLACE FUNCTION public.create_ops_daily_shadow_partitions(p_day_start timestamptz) +RETURNS void +LANGUAGE plpgsql +SECURITY INVOKER +SET search_path = pg_catalog, public +SET "TimeZone" = 'UTC' +SET lock_timeout = '2s' +AS $$ +DECLARE + day_end timestamptz; + day_suffix text; + parent_name text; + parent_oid oid; + partition_name text; + partition_oid oid; + attached_parent_oid oid; + partition_bound text; + utc_start_literal text; + utc_end_literal text; +BEGIN + IF p_day_start IS NULL THEN + RAISE EXCEPTION 'p_day_start is required'; + END IF; + + IF p_day_start <> ( + pg_catalog.date_trunc('day', p_day_start AT TIME ZONE 'UTC') AT TIME ZONE 'UTC' + ) THEN + RAISE EXCEPTION 'p_day_start must be an exact UTC day boundary: %', p_day_start; + END IF; + + IF NOT pg_catalog.pg_try_advisory_xact_lock( + pg_catalog.hashtextextended('ops_daily_shadow_partition_manager', 0) + ) THEN + RAISE EXCEPTION 'another OPS shadow partition operation is in progress'; + END IF; + + day_end := p_day_start + INTERVAL '1 day'; + day_suffix := pg_catalog.to_char(p_day_start AT TIME ZONE 'UTC', 'YYYYMMDD'); + utc_start_literal := pg_catalog.to_char( + p_day_start AT TIME ZONE 'UTC', + 'YYYY-MM-DD"T"HH24:MI:SS"Z"' + ); + utc_end_literal := pg_catalog.to_char( + day_end AT TIME ZONE 'UTC', + 'YYYY-MM-DD"T"HH24:MI:SS"Z"' + ); + + FOREACH parent_name IN ARRAY ARRAY[ + 'ops_system_logs_daily_shadow', + 'ops_error_logs_daily_shadow' + ] + LOOP + SELECT c.oid + INTO parent_oid + FROM pg_catalog.pg_class AS c + JOIN pg_catalog.pg_namespace AS n ON n.oid = c.relnamespace + WHERE n.nspname = 'public' + AND c.relname = parent_name + AND c.relkind = 'p'; + + IF parent_oid IS NULL THEN + RAISE EXCEPTION 'required partitioned parent public.% does not exist', parent_name; + END IF; + + partition_name := parent_name || '_' || day_suffix; + EXECUTE pg_catalog.format( + 'CREATE TABLE IF NOT EXISTS public.%I PARTITION OF public.%I ' + 'FOR VALUES FROM (%L::timestamptz) TO (%L::timestamptz)', + partition_name, + parent_name, + utc_start_literal, + utc_end_literal + ); + + SELECT c.oid, + i.inhparent, + pg_catalog.pg_get_expr(c.relpartbound, c.oid, true) + INTO partition_oid, attached_parent_oid, partition_bound + FROM pg_catalog.pg_class AS c + JOIN pg_catalog.pg_namespace AS n ON n.oid = c.relnamespace + LEFT JOIN pg_catalog.pg_inherits AS i ON i.inhrelid = c.oid + WHERE n.nspname = 'public' + AND c.relname = partition_name + AND c.relispartition; + + IF partition_oid IS NULL OR attached_parent_oid IS DISTINCT FROM parent_oid THEN + RAISE EXCEPTION 'public.% exists but is not attached to public.%', partition_name, parent_name; + END IF; + + IF pg_catalog.strpos( + partition_bound, + pg_catalog.to_char(p_day_start AT TIME ZONE 'UTC', 'YYYY-MM-DD HH24:MI:SS') + ) = 0 OR pg_catalog.strpos( + partition_bound, + pg_catalog.to_char(day_end AT TIME ZONE 'UTC', 'YYYY-MM-DD HH24:MI:SS') + ) = 0 THEN + RAISE EXCEPTION 'public.% has unexpected bounds: %', partition_name, partition_bound; + END IF; + END LOOP; +END; +$$; + +REVOKE ALL ON FUNCTION public.create_ops_daily_shadow_partitions(timestamptz) FROM PUBLIC; + +COMMENT ON FUNCTION public.create_ops_daily_shadow_partitions(timestamptz) IS + 'Explicitly creates one UTC day partition for both empty OPS shadow parents. It never copies, attaches, renames, or switches production tables.'; diff --git a/backend/migrations/216_usage_log_image_input_tokens.sql b/backend/migrations/216_usage_log_image_input_tokens.sql new file mode 100644 index 000000000..0b05acb99 --- /dev/null +++ b/backend/migrations/216_usage_log_image_input_tokens.sql @@ -0,0 +1,9 @@ +-- 216_usage_log_image_input_tokens.sql +-- usage_logs 单独记录图片输入 token 与费用，便于图片编辑、图生图等场景对账。 +-- input_tokens 继续保留总输入 token；input_cost 改为仅记录文本输入费用， +-- image_input_cost 单独记录图片输入费用，total_cost 与 actual_cost 总额口径不变。 +SET LOCAL lock_timeout = '2s'; +SET LOCAL statement_timeout = '30s'; + +ALTER TABLE usage_logs ADD COLUMN IF NOT EXISTS image_input_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE usage_logs ADD COLUMN IF NOT EXISTS image_input_cost DECIMAL(20, 10) NOT NULL DEFAULT 0; diff --git a/backend/migrations/217_openai_owned_agent_identity_unique_notx.sql b/backend/migrations/217_openai_owned_agent_identity_unique_notx.sql new file mode 100644 index 000000000..e5dfae03f --- /dev/null +++ b/backend/migrations/217_openai_owned_agent_identity_unique_notx.sql @@ -0,0 +1,75 @@ +CREATE UNIQUE INDEX CONCURRENTLY IF NOT EXISTS idx_accounts_owned_openai_org_user_v2_uniq + ON public.accounts ( + owner_user_id, + LOWER(NULLIF(BTRIM(credentials->>'organization_id'), '')), + NULLIF(BTRIM(credentials->>'chatgpt_user_id'), '') + ) + WHERE deleted_at IS NULL + AND owner_user_id IS NOT NULL + AND platform = 'openai' + AND type = 'oauth' + AND COALESCE(LOWER(NULLIF(BTRIM(credentials->>'auth_mode'), '')), '') <> 'agentidentity' + AND NULLIF(BTRIM(credentials->>'organization_id'), '') IS NOT NULL + AND NULLIF(BTRIM(credentials->>'chatgpt_user_id'), '') IS NOT NULL; + +CREATE UNIQUE INDEX CONCURRENTLY IF NOT EXISTS idx_accounts_owned_openai_org_account_v2_uniq + ON public.accounts ( + owner_user_id, + LOWER(NULLIF(BTRIM(credentials->>'organization_id'), '')), + NULLIF(BTRIM(credentials->>'chatgpt_account_id'), '') + ) + WHERE deleted_at IS NULL + AND owner_user_id IS NOT NULL + AND platform = 'openai' + AND type = 'oauth' + AND COALESCE(LOWER(NULLIF(BTRIM(credentials->>'auth_mode'), '')), '') <> 'agentidentity' + AND NULLIF(BTRIM(credentials->>'organization_id'), '') IS NOT NULL + AND NULLIF(BTRIM(credentials->>'chatgpt_user_id'), '') IS NULL + AND NULLIF(BTRIM(credentials->>'chatgpt_account_id'), '') IS NOT NULL; + +CREATE UNIQUE INDEX CONCURRENTLY IF NOT EXISTS idx_accounts_owned_openai_legacy_user_v2_uniq + ON public.accounts ( + owner_user_id, + NULLIF(BTRIM(credentials->>'chatgpt_user_id'), '') + ) + WHERE deleted_at IS NULL + AND owner_user_id IS NOT NULL + AND platform = 'openai' + AND type = 'oauth' + AND COALESCE(LOWER(NULLIF(BTRIM(credentials->>'auth_mode'), '')), '') <> 'agentidentity' + AND NULLIF(BTRIM(credentials->>'organization_id'), '') IS NULL + AND NULLIF(BTRIM(credentials->>'chatgpt_user_id'), '') IS NOT NULL; + +CREATE UNIQUE INDEX CONCURRENTLY IF NOT EXISTS idx_accounts_owned_openai_legacy_account_v2_uniq + ON public.accounts ( + owner_user_id, + NULLIF(BTRIM(credentials->>'chatgpt_account_id'), '') + ) + WHERE deleted_at IS NULL + AND owner_user_id IS NOT NULL + AND platform = 'openai' + AND type = 'oauth' + AND COALESCE(LOWER(NULLIF(BTRIM(credentials->>'auth_mode'), '')), '') <> 'agentidentity' + AND NULLIF(BTRIM(credentials->>'organization_id'), '') IS NULL + AND NULLIF(BTRIM(credentials->>'chatgpt_user_id'), '') IS NULL + AND NULLIF(BTRIM(credentials->>'chatgpt_account_id'), '') IS NOT NULL; + +CREATE UNIQUE INDEX CONCURRENTLY IF NOT EXISTS idx_accounts_owned_openai_agent_identity_team_uniq + ON public.accounts ( + owner_user_id, + NULLIF(BTRIM(credentials->>'chatgpt_account_id'), '') + ) + WHERE deleted_at IS NULL + AND owner_user_id IS NOT NULL + AND platform = 'openai' + AND type = 'oauth' + AND LOWER(NULLIF(BTRIM(credentials->>'auth_mode'), '')) = 'agentidentity' + AND NULLIF(BTRIM(credentials->>'chatgpt_account_id'), '') IS NOT NULL; + +DROP INDEX CONCURRENTLY IF EXISTS public.idx_accounts_owned_openai_org_user_uniq; + +DROP INDEX CONCURRENTLY IF EXISTS public.idx_accounts_owned_openai_org_account_uniq; + +DROP INDEX CONCURRENTLY IF EXISTS public.idx_accounts_owned_openai_legacy_user_uniq; + +DROP INDEX CONCURRENTLY IF EXISTS public.idx_accounts_owned_openai_legacy_account_uniq; diff --git a/backend/migrations/openai_owned_agent_identity_unique_test.go b/backend/migrations/openai_owned_agent_identity_unique_test.go new file mode 100644 index 000000000..c76a45363 --- /dev/null +++ b/backend/migrations/openai_owned_agent_identity_unique_test.go @@ -0,0 +1,69 @@ +package migrations + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +const openAIOwnedAgentIdentityUniqueMigration = "217_openai_owned_agent_identity_unique_notx.sql" + +func TestOpenAIOwnedAgentIdentityUniqueMigrationBuildsBeforeDropping(t *testing.T) { + content, err := FS.ReadFile(openAIOwnedAgentIdentityUniqueMigration) + require.NoError(t, err) + + sql := string(content) + require.Equal(t, 5, strings.Count(sql, "CREATE UNIQUE INDEX CONCURRENTLY IF NOT EXISTS")) + require.Equal(t, 4, strings.Count(sql, "DROP INDEX CONCURRENTLY IF EXISTS")) + require.Equal(t, 5, strings.Count(sql, "ON public.accounts (")) + + firstDrop := strings.Index(sql, "DROP INDEX CONCURRENTLY IF EXISTS") + require.Positive(t, firstDrop) + for _, indexName := range []string{ + "idx_accounts_owned_openai_org_user_v2_uniq", + "idx_accounts_owned_openai_org_account_v2_uniq", + "idx_accounts_owned_openai_legacy_user_v2_uniq", + "idx_accounts_owned_openai_legacy_account_v2_uniq", + "idx_accounts_owned_openai_agent_identity_team_uniq", + } { + position := strings.Index(sql, "CREATE UNIQUE INDEX CONCURRENTLY IF NOT EXISTS "+indexName) + require.GreaterOrEqual(t, position, 0, indexName) + require.Less(t, position, firstDrop, indexName) + } + for _, indexName := range []string{ + "idx_accounts_owned_openai_org_user_uniq", + "idx_accounts_owned_openai_org_account_uniq", + "idx_accounts_owned_openai_legacy_user_uniq", + "idx_accounts_owned_openai_legacy_account_uniq", + } { + require.Contains(t, sql, "DROP INDEX CONCURRENTLY IF EXISTS public."+indexName) + } + require.NotContains(t, sql[firstDrop:], "CREATE UNIQUE INDEX") +} + +func TestOpenAIOwnedAgentIdentityUniqueMigrationSeparatesOAuthAndTeamSemantics(t *testing.T) { + content, err := FS.ReadFile(openAIOwnedAgentIdentityUniqueMigration) + require.NoError(t, err) + sql := string(content) + + nonAgentPredicate := "COALESCE(LOWER(NULLIF(BTRIM(credentials->>'auth_mode'), '')), '') <> 'agentidentity'" + require.Equal(t, 4, strings.Count(sql, nonAgentPredicate)) + + teamStart := strings.Index(sql, "CREATE UNIQUE INDEX CONCURRENTLY IF NOT EXISTS idx_accounts_owned_openai_agent_identity_team_uniq") + firstDrop := strings.Index(sql, "DROP INDEX CONCURRENTLY IF EXISTS") + require.Positive(t, teamStart) + require.Greater(t, firstDrop, teamStart) + + teamSQL := sql[teamStart:firstDrop] + require.Contains(t, teamSQL, "owner_user_id") + require.Contains(t, teamSQL, "credentials->>'chatgpt_account_id'") + require.NotContains(t, teamSQL, "chatgpt_user_id") + require.Contains(t, teamSQL, "LOWER(NULLIF(BTRIM(credentials->>'auth_mode'), '')) = 'agentidentity'") + require.NotContains(t, sql, "idx_accounts_owned_openai_agent_identity_runtime_uniq") + + upperSQL := strings.ToUpper(sql) + require.NotContains(t, upperSQL, "UPDATE ACCOUNTS") + require.NotContains(t, upperSQL, "DELETE FROM ACCOUNTS") + require.NotContains(t, upperSQL, "INSERT INTO ACCOUNTS") +} diff --git a/backend/migrations/ops_daily_partition_shadow_test.go b/backend/migrations/ops_daily_partition_shadow_test.go new file mode 100644 index 000000000..d0ac90346 --- /dev/null +++ b/backend/migrations/ops_daily_partition_shadow_test.go @@ -0,0 +1,62 @@ +package migrations + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +const opsDailyPartitionShadowMigration = "215_ops_daily_partition_shadow.sql" + +func readNormalizedOpsDailyPartitionShadowMigration(t *testing.T) string { + t.Helper() + + content, err := FS.ReadFile(opsDailyPartitionShadowMigration) + require.NoError(t, err) + return strings.Join(strings.Fields(string(content)), " ") +} + +func TestOpsDailyPartitionShadowMigrationCreatesOnlyEmptyParents(t *testing.T) { + sql := readNormalizedOpsDailyPartitionShadowMigration(t) + + require.Contains(t, sql, "CREATE TABLE IF NOT EXISTS public.ops_system_logs_daily_shadow") + require.Contains(t, sql, "CREATE TABLE IF NOT EXISTS public.ops_error_logs_daily_shadow") + require.Contains(t, sql, "PARTITION BY RANGE (created_at)") + require.Contains(t, sql, "PRIMARY KEY (created_at, id)") + require.Contains(t, sql, "id_index_name := parent_name || '_id_idx'") + require.Contains(t, sql, "'CREATE INDEX %I ON public.%I (id)'") + require.Contains(t, sql, "ALTER TABLE public.ops_system_logs_daily_shadow ALTER COLUMN id DROP DEFAULT") + require.Contains(t, sql, "ALTER TABLE public.ops_error_logs_daily_shadow ALTER COLUMN id DROP DEFAULT") + + upperSQL := strings.ToUpper(sql) + require.NotContains(t, upperSQL, "ALTER TABLE PUBLIC.OPS_SYSTEM_LOGS RENAME") + require.NotContains(t, upperSQL, "ALTER TABLE PUBLIC.OPS_ERROR_LOGS RENAME") + require.NotContains(t, upperSQL, "ATTACH PARTITION") + require.NotContains(t, upperSQL, "INSERT INTO PUBLIC.OPS_SYSTEM_LOGS_DAILY_SHADOW") + require.NotContains(t, upperSQL, "INSERT INTO PUBLIC.OPS_ERROR_LOGS_DAILY_SHADOW") + require.NotContains(t, upperSQL, "CREATE TABLE AS") + require.NotContains(t, upperSQL, "COPY ") +} + +func TestOpsDailyPartitionShadowMigrationRequiresExplicitUTCDay(t *testing.T) { + sql := readNormalizedOpsDailyPartitionShadowMigration(t) + + require.Contains(t, sql, "CREATE OR REPLACE FUNCTION public.create_ops_daily_shadow_partitions(p_day_start timestamptz)") + require.Contains(t, sql, "SET \"TimeZone\" = 'UTC'") + require.Contains(t, sql, "p_day_start must be an exact UTC day boundary") + require.Contains(t, sql, `YYYY-MM-DD"T"HH24:MI:SS"Z"`) + require.Contains(t, sql, "FOR VALUES FROM (%L::timestamptz) TO (%L::timestamptz)") + require.Contains(t, sql, "REVOKE ALL ON FUNCTION public.create_ops_daily_shadow_partitions(timestamptz) FROM PUBLIC") +} + +func TestOpsDailyPartitionShadowMigrationDoesNotInvokePartitionCreation(t *testing.T) { + content, err := FS.ReadFile(opsDailyPartitionShadowMigration) + require.NoError(t, err) + + sql := string(content) + require.Equal(t, 3, strings.Count(sql, "create_ops_daily_shadow_partitions"), + "the helper may only appear in its declaration and privilege/comment signature") + require.NotContains(t, sql, "SELECT public.create_ops_daily_shadow_partitions(") + require.NotContains(t, sql, "PERFORM public.create_ops_daily_shadow_partitions(") +} diff --git a/deploy/.env.example b/deploy/.env.example index b1eb3f121..340dd041c 100644 --- a/deploy/.env.example +++ b/deploy/.env.example @@ -20,9 +20,18 @@ SERVER_PORT=8080 # Server mode: release or debug SERVER_MODE=release +# 可选：仅用于 GitHub Release API 更新检查，下载发布资产时不会携带该 Token。 +UPDATE_GITHUB_TOKEN= + +# Pixel 默认禁止用上游原版二进制覆盖当前二开版本；仅明确接受风险时启用。 +PIXEL_ALLOW_UPSTREAM_IN_PLACE_UPDATE=false + # 为已认证的管理端和用户端 Web API 返回 Server-Timing；默认关闭 ENABLE_SERVER_TIMING=false +# HTTP 服务优雅退出与应用清理各阶段的独立超时（秒） +SERVER_SHUTDOWN_TIMEOUT_SECONDS=30 + # ----------------------------------------------------------------------------- # Logging Configuration # 日志配置 diff --git a/deploy/config.example.yaml b/deploy/config.example.yaml index a52585b19..8fe4aa589 100644 --- a/deploy/config.example.yaml +++ b/deploy/config.example.yaml @@ -27,6 +27,9 @@ server: # 用于生成邮件中的外部链接（例如：重置密码链接）的前端基础地址 # Example: "https://example.com" frontend_url: "" + # Per-phase timeout for graceful HTTP shutdown and application cleanup (seconds) + # HTTP 服务优雅退出与应用清理各阶段的独立超时（秒） + shutdown_timeout_seconds: 30 # Trusted proxies for X-Forwarded-For parsing (CIDR/IP). Empty disables trusted proxies. # 信任的代理地址（CIDR/IP 格式），用于解析 X-Forwarded-For 头。留空则禁用代理信任。 trusted_proxies: [] @@ -86,6 +89,13 @@ cors: # 安全配置 # ============================================================================= security: + # Additional headers that may carry a client IP. They are only read when the + # direct peer matches server.trusted_proxies; direct clients cannot spoof them. + # 可携带客户端 IP 的自定义请求头。仅当直连来源命中 server.trusted_proxies 时才会解析， + # 直连客户端无法通过伪造这些请求头覆盖自身 IP。 + # 可信代理必须覆盖（而不是透传）同名客户端请求头，否则代理后的客户端仍可能伪造该值。 + # Example / 示例: ["True-Client-IP", "X-CDN-Client-IP"] + forwarded_client_ip_headers: [] url_allowlist: # Enable URL allowlist validation (disable to skip all URL checks) # 启用 URL 白名单验证（禁用则跳过所有 URL 检查） @@ -239,12 +249,13 @@ gateway: # 默认 false：过滤超时头，降低上游提前断流风险。 openai_passthrough_allow_timeout_headers: false # OpenAI 原生 HTTP Responses 首个语义输出超时（秒，0=禁用）。 - # 超时从网关收到请求开始计算，包含等待响应头；不适用于 passthrough 或 WebSocket 链路。 + # 默认 60 秒；从网关完成请求解析并进入调度阶段开始计算，包含选号/排队后的剩余预算和等待响应头； + # 只保护首个语义事件，已开始输出的长流不受此值截断。适用于原生及 passthrough HTTP Responses，不适用于 WebSocket 链路。 # 首次输出前单次尝试最多暂存 8 MiB，超限时切换账号且不暴露不完整 SSE 数据。 # 注意：超时请求可能已产生上游用量，切换账号重试可能导致上游重复计费。 - openai_first_output_timeout_seconds: 0 - # high/xhigh/max 推理可选覆盖值（秒，0=使用普通首输出超时）。 - openai_high_effort_first_output_timeout_seconds: 0 + openai_first_output_timeout_seconds: 60 + # high/xhigh/max 推理可选覆盖值（秒，0=使用普通首输出超时）；默认 180 秒。 + openai_high_effort_first_output_timeout_seconds: 180 # OpenAI Responses WebSocket 配置（默认开启，可按需回滚到 HTTP） openai_ws: # 新版 WS mode 路由（默认关闭）。关闭时保持当前 legacy 实现行为。 @@ -846,13 +857,44 @@ ops: # 是否启用运维监控功能（后台任务和接口） # Set to false to hide ops menu in sidebar and disable all ops features # 设置为 false 可在左侧栏隐藏运维监控菜单并禁用所有运维监控功能 - # Other detailed settings (cleanup, aggregation, etc.) are configured in ops settings dialog - # 其他详细设置（数据清理、预聚合等）在运维监控设置对话框中配置 + # This file is the hard safety boundary. The ops dialog overrides cleanup enabled/schedule/retention at runtime; + # window, timeout, batch and archive-expiry controls below remain static and take effect after restart. + # 本文件是硬安全边界。运维设置对话框会动态覆盖清理开关、计划和保留天数； + # 窗口、超时、批次和归档过期参数仍由本文件控制，并在重启后生效。 enabled: true cleanup: + # Hard switch for the ops cleanup worker. The runtime dialog cannot enable cleanup when this is false. + # 运维清理任务硬开关；此处为 false 时，运行时设置无法启用清理。 + enabled: true + # Fallback schedule when no persisted advanced setting exists. + # 没有持久化高级设置时使用的回退计划。 + schedule: "0 4 * * *" # Expiration days for archived ops_error_logs / ops_system_logs objects. 0 means never expire. # ops_error_logs / ops_system_logs 归档对象过期天数，0 表示永不过期。 - archive_expire_days: 14 + archive_expire_days: 30 + # Calendar-day span of each archive/delete window. + # 每个归档并删除窗口覆盖的自然日数。 + archive_window_days: 1 + # Maximum windows advanced per log table in one scheduled run. + # 每次任务对每张日志表最多追赶的窗口数。 + max_catchup_windows_per_run: 2 + # Timeout for one table/window archive and upload. + # 单张表单个窗口归档上传超时（秒）。 + archive_timeout_seconds: 1800 + # Timeout for one table/window or one auxiliary-table deletion phase. + # 单张表单个窗口或辅助表删除阶段超时（秒）。 + delete_timeout_seconds: 1800 + # Total timeout for one scheduled cleanup run. + # 单次定时清理总超时（秒）。 + run_timeout_seconds: 18000 + # Rows deleted per SQL statement. + # 单条 SQL 删除批次大小。 + delete_batch_size: 5000 + # Fallback retention days when no persisted advanced setting exists. 0 disables that target. + # 没有持久化高级设置时使用的回退保留天数；0 表示禁用对应目标，不会清空表。 + error_log_retention_days: 30 + minute_metrics_retention_days: 30 + hourly_metrics_retention_days: 30 # ============================================================================= # JWT Configuration diff --git a/deploy/docker-compose.dev.yml b/deploy/docker-compose.dev.yml index 7793e424a..38de0d768 100644 --- a/deploy/docker-compose.dev.yml +++ b/deploy/docker-compose.dev.yml @@ -83,10 +83,10 @@ services: - ./redis_data:/data command: > sh -c ' - redis-server - --save 60 1 - --appendonly yes - --appendfsync everysec + redis-server \ + --save 60 1 \ + --appendonly yes \ + --appendfsync everysec \ ${REDIS_PASSWORD:+--requirepass "$REDIS_PASSWORD"}' environment: - TZ=${TZ:-Asia/Shanghai} diff --git a/deploy/docker-compose.local.yml b/deploy/docker-compose.local.yml index 5aea78fb2..9aca92107 100644 --- a/deploy/docker-compose.local.yml +++ b/deploy/docker-compose.local.yml @@ -52,6 +52,8 @@ services: - SERVER_PORT=8080 - SERVER_MODE=${SERVER_MODE:-release} - RUN_MODE=${RUN_MODE:-standard} + - UPDATE_GITHUB_TOKEN=${UPDATE_GITHUB_TOKEN:-} + - PIXEL_ALLOW_UPSTREAM_IN_PLACE_UPDATE=${PIXEL_ALLOW_UPSTREAM_IN_PLACE_UPDATE:-false} # ======================================================================= # Database Configuration (PostgreSQL) @@ -207,10 +209,10 @@ services: - ./redis_data:/data command: > sh -c ' - redis-server - --save 60 1 - --appendonly yes - --appendfsync everysec + redis-server \ + --save 60 1 \ + --appendonly yes \ + --appendfsync everysec \ ${REDIS_PASSWORD:+--requirepass "$REDIS_PASSWORD"}' environment: - TZ=${TZ:-Asia/Shanghai} diff --git a/deploy/docker-compose.yml b/deploy/docker-compose.yml index b7335bd87..9db5adae5 100644 --- a/deploy/docker-compose.yml +++ b/deploy/docker-compose.yml @@ -48,6 +48,8 @@ services: - SERVER_PORT=8080 - SERVER_MODE=${SERVER_MODE:-release} - RUN_MODE=${RUN_MODE:-standard} + - UPDATE_GITHUB_TOKEN=${UPDATE_GITHUB_TOKEN:-} + - PIXEL_ALLOW_UPSTREAM_IN_PLACE_UPDATE=${PIXEL_ALLOW_UPSTREAM_IN_PLACE_UPDATE:-false} # ======================================================================= # Database Configuration (PostgreSQL) @@ -192,6 +194,12 @@ services: nofile: soft: 100000 hard: 100000 + command: > + postgres + -c max_connections=${POSTGRES_MAX_CONNECTIONS:-100} + -c shared_buffers=${POSTGRES_SHARED_BUFFERS:-128MB} + -c effective_cache_size=${POSTGRES_EFFECTIVE_CACHE_SIZE:-4GB} + -c maintenance_work_mem=${POSTGRES_MAINTENANCE_WORK_MEM:-64MB} volumes: - postgres_data:/var/lib/postgresql/data environment: @@ -233,10 +241,10 @@ services: - redis_data:/data command: > sh -c ' - redis-server - --save 60 1 - --appendonly yes - --appendfsync everysec + redis-server \ + --save 60 1 \ + --appendonly yes \ + --appendfsync everysec \ --requirepass "$$REDIS_PASSWORD"' environment: - TZ=${TZ:-Asia/Shanghai} diff --git a/deploy/pixel-retention-cleanup.service b/deploy/pixel-retention-cleanup.service new file mode 100644 index 000000000..0b1674179 --- /dev/null +++ b/deploy/pixel-retention-cleanup.service @@ -0,0 +1,18 @@ +[Unit] +Description=Remove expired Pixel release and upload artifacts +After=local-fs.target pixel.service +ConditionPathIsDirectory=/opt/sub2api/releases +ConditionPathIsDirectory=/home/pixel/sub2api_release_uploads + +[Service] +Type=oneshot +ExecStart=/usr/local/sbin/pixel-retention-cleanup +User=root +Group=root +Nice=10 +IOSchedulingClass=idle +NoNewPrivileges=true +PrivateTmp=true +ProtectSystem=full +ProtectHome=read-only +ReadWritePaths=/opt/sub2api/releases /home/pixel/sub2api_release_uploads /run diff --git a/deploy/pixel-retention-cleanup.sh b/deploy/pixel-retention-cleanup.sh new file mode 100644 index 000000000..c9babb21d --- /dev/null +++ b/deploy/pixel-retention-cleanup.sh @@ -0,0 +1,101 @@ +#!/usr/bin/env bash + +set -euo pipefail + +readonly retention_minutes=43200 +readonly release_root="/opt/sub2api/releases" +readonly upload_root="/home/pixel/sub2api_release_uploads" +readonly lock_file="/run/pixel-release-maintenance.lock" + +exec 9>"${lock_file}" +if ! flock -n 9; then + echo "retention cleanup is already running" >&2 + exit 1 +fi + +resolved_release_root=$(readlink -f -- "${release_root}") +resolved_upload_root=$(readlink -f -- "${upload_root}") +current_release=$(readlink -f -- /opt/sub2api/current) + +if [[ "${resolved_release_root}" != "${release_root}" ]]; then + echo "unexpected release root: ${resolved_release_root}" >&2 + exit 1 +fi +if [[ "${resolved_upload_root}" != "${upload_root}" ]]; then + echo "unexpected upload root: ${resolved_upload_root}" >&2 + exit 1 +fi +deleted_releases=0 +deleted_uploads=0 + +while IFS= read -r -d '' candidate; do + resolved_candidate=$(readlink -f -- "${candidate}") + case "${resolved_candidate}" in + "${resolved_release_root}"/*) ;; + *) + echo "unsafe release path: ${resolved_candidate}" >&2 + exit 1 + ;; + esac + + # Re-read all protection pointers immediately before deletion. The deployment + # switch uses the same flock so rollback/current cannot change inside this check. + current_release=$(readlink -f -- /opt/sub2api/current) + pixel_pid=$(systemctl show pixel.service -p MainPID --value) + if [[ -z "${pixel_pid}" || "${pixel_pid}" == "0" ]]; then + echo "pixel.service is not running; skipping release deletion" + continue + fi + pixel_executable=$(readlink -f -- "/proc/${pixel_pid}/exe" || true) + pixel_working_directory=$(readlink -f -- "/proc/${pixel_pid}/cwd" || true) + + if [[ "${resolved_candidate}" == "${current_release}" ]]; then + echo "skipping current release: ${resolved_candidate}" + continue + fi + case "${pixel_executable}" in + "${resolved_candidate}"/*) + echo "skipping active executable release: ${resolved_candidate}" + continue + ;; + esac + case "${pixel_working_directory}" in + "${resolved_candidate}"|"${resolved_candidate}"/*) + echo "skipping active working directory: ${resolved_candidate}" + continue + ;; + esac + if mountpoint -q -- "${resolved_candidate}"; then + echo "skipping mounted release candidate: ${resolved_candidate}" + continue + fi + + rm -rf --one-file-system -- "${resolved_candidate}" + deleted_releases=$((deleted_releases + 1)) +done < <( + find "${resolved_release_root}" -xdev -mindepth 1 -maxdepth 1 -type d \ + -mmin "+${retention_minutes}" -print0 +) + +while IFS= read -r -d '' candidate; do + resolved_candidate=$(readlink -f -- "${candidate}") + case "${resolved_candidate}" in + "${resolved_upload_root}"/*) ;; + *) + echo "unsafe upload path: ${resolved_candidate}" >&2 + exit 1 + ;; + esac + + rm -f -- "${resolved_candidate}" + deleted_uploads=$((deleted_uploads + 1)) +done < <( + find "${resolved_upload_root}" -mindepth 1 -maxdepth 1 -type f \ + -mmin "+${retention_minutes}" -print0 +) + +find "${resolved_upload_root}" -mindepth 1 -maxdepth 1 -type d -empty \ + -mmin "+${retention_minutes}" -delete + +printf 'retention cleanup completed: releases=%d uploads=%d current=%s\n' \ + "${deleted_releases}" "${deleted_uploads}" "${current_release}" diff --git a/deploy/pixel-retention-cleanup.timer b/deploy/pixel-retention-cleanup.timer new file mode 100644 index 000000000..74f42052d --- /dev/null +++ b/deploy/pixel-retention-cleanup.timer @@ -0,0 +1,11 @@ +[Unit] +Description=Run Pixel artifact retention cleanup daily + +[Timer] +OnCalendar=*-*-* 03:15:00 Asia/Shanghai +RandomizedDelaySec=5m +AccuracySec=1m +Unit=pixel-retention-cleanup.service + +[Install] +WantedBy=timers.target diff --git a/deploy/rsyslog.logrotate b/deploy/rsyslog.logrotate new file mode 100644 index 000000000..18f69a8d4 --- /dev/null +++ b/deploy/rsyslog.logrotate @@ -0,0 +1,18 @@ +/var/log/cron +/var/log/maillog +/var/log/messages +/var/log/secure +/var/log/spooler +{ + daily + rotate 30 + maxage 30 + compress + delaycompress + missingok + notifempty + sharedscripts + postrotate + /usr/bin/systemctl -s HUP kill rsyslog.service >/dev/null 2>&1 || true + endscript +} diff --git a/frontend/src/App.vue b/frontend/src/App.vue index 7485aa1af..aeea1fbc9 100644 --- a/frontend/src/App.vue +++ b/frontend/src/App.vue @@ -1,12 +1,14 @@ diff --git a/frontend/src/components/common/DataTable.vue b/frontend/src/components/common/DataTable.vue index f7c170c7f..5bf794ad7 100644 --- a/frontend/src/components/common/DataTable.vue +++ b/frontend/src/components/common/DataTable.vue @@ -1,7 +1,51 @@