diff --git a/Dockerfile b/Dockerfile index 17b631d5a..b2dfd7034 100644 --- a/Dockerfile +++ b/Dockerfile @@ -18,7 +18,9 @@ ARG NPM_CONFIG_REGISTRY= # ----------------------------------------------------------------------------- # Stage 1: Frontend Builder # ----------------------------------------------------------------------------- -FROM ${NODE_IMAGE} AS frontend-builder +# --platform=$BUILDPLATFORM: the frontend output is JS (arch-neutral), so build +# it on the native host arch instead of under QEMU emulation for the target. +FROM --platform=${BUILDPLATFORM} ${NODE_IMAGE} AS frontend-builder ARG NPM_CONFIG_REGISTRY WORKDIR /app/frontend @@ -44,7 +46,11 @@ RUN pnpm run build # ----------------------------------------------------------------------------- # Stage 2: Backend Builder # ----------------------------------------------------------------------------- -FROM ${GOLANG_IMAGE} AS backend-builder +# --platform=$BUILDPLATFORM: run the Go toolchain on the native host arch and +# cross-compile to the target arch below. The binary is CGO_ENABLED=0, so this +# is a clean pure-Go cross-compile — no QEMU emulation of go mod download / go +# build (emulated networking here was dropping module fetches with EOF). +FROM --platform=${BUILDPLATFORM} ${GOLANG_IMAGE} AS backend-builder # Build arguments for version info (set by CI) ARG VERSION= @@ -52,6 +58,9 @@ ARG COMMIT=docker ARG DATE ARG GOPROXY ARG GOSUMDB +# Populated by buildx from the --platform target (e.g. linux/amd64). +ARG TARGETOS +ARG TARGETARCH ENV GOPROXY=${GOPROXY} ENV GOSUMDB=${GOSUMDB} @@ -63,7 +72,10 @@ WORKDIR /app/backend # Copy go mod files first (better caching) COPY backend/go.mod backend/go.sum ./ -RUN go mod download +# Cache mount keeps the module cache across builds so a transient CDN blip on +# retry resumes instead of re-fetching every zip from scratch. +RUN --mount=type=cache,id=sub2api-gomod,target=/go/pkg/mod \ + go mod download # Copy backend source first COPY backend/ ./ @@ -73,10 +85,12 @@ COPY --from=frontend-builder /app/backend/internal/web/dist ./internal/web/dist # Build the binary (BuildType=release for CI builds, embed frontend) # Version precedence: build arg VERSION > exact git tag > cmd/server/VERSION -RUN VERSION_VALUE="${VERSION}" && \ +RUN --mount=type=cache,id=sub2api-gomod,target=/go/pkg/mod \ + --mount=type=cache,id=sub2api-gobuild,target=/root/.cache/go-build \ + VERSION_VALUE="${VERSION}" && \ if [ -z "${VERSION_VALUE}" ]; then VERSION_VALUE="$(./scripts/resolve-version.sh)"; fi && \ DATE_VALUE="${DATE:-$(date -u +%Y-%m-%dT%H:%M:%SZ)}" && \ - CGO_ENABLED=0 GOOS=linux go build \ + CGO_ENABLED=0 GOOS=${TARGETOS:-linux} GOARCH=${TARGETARCH} go build \ -tags embed \ -ldflags="-s -w -X main.Version=${VERSION_VALUE} -X main.Commit=${COMMIT} -X main.Date=${DATE_VALUE} -X main.BuildType=release" \ -trimpath \ diff --git a/README.md b/README.md index dd415c307..f48e09087 100644 --- a/README.md +++ b/README.md @@ -639,8 +639,20 @@ override this limit. The connection cap is coordinated through Redis using a 60-second lease that is refreshed every 20 seconds. A process that cannot confirm a lease for a full lease lifetime closes its local WebSocket rather than continuing outside -the global cap. Use `http_bridge` for client-WebSocket/upstream-HTTP operation -when rolling out or mitigating upstream WebSocket issues. +the global cap. + +Enable the v2 mode router before selecting an account-level WS mode such as +`http_bridge`: + +```yaml +gateway: + openai_ws: + mode_router_v2_enabled: true +``` + +Or set `GATEWAY_OPENAI_WS_MODE_ROUTER_V2_ENABLED=true` in the environment. +Use `http_bridge` for client-WebSocket/upstream-HTTP operation when rolling out +or mitigating upstream WebSocket issues. #### ⚠️ Important: Creating the Admin Account @@ -788,9 +800,9 @@ xAI quota is passive. Sub2API does not invent subscription quota values; it reco `401` responses temporarily remove accounts with invalid credentials from scheduling. `403` responses are treated as access or entitlement failures instead of token-refresh loops. `429` responses use `Retry-After` or a short cooldown to temporarily remove the account from scheduling. -New Grok image and video generation requests use a media-specific eligibility check. An OAuth account is excluded from new media generation when its recorded weekly or monthly billing probe returns `403`; chat requests and video status lookups are not affected by this media-only quarantine. If no eligible account remains, the media endpoint returns HTTP `503` with error type `grok_media_no_eligible_account` instead of forwarding the request to a known-ineligible account. +New Grok image and video generation requests use a media-specific eligibility check. API-key accounts remain eligible. OAuth accounts require positive paid-entitlement evidence from the xAI billing probe; Free, forbidden, missing, malformed, and inconclusive billing observations are excluded from new media generation. Unobserved OAuth accounts are probed before the first media request is forwarded, and imports run the billing-first quota probe proactively. Chat requests and video status lookups are not affected by this media-only quarantine. If no eligible account remains, the media endpoint returns HTTP `503` with error type `grok_media_no_eligible_account`. -Administrators can override automatic media eligibility through the account create/update API by setting `extra.grok_media_eligible` to `false` (exclude) or `true` (force eligible). On update, set it to `null` to remove the override and return to automatic probe-based behavior; omitting the field preserves the current override. A missing billing observation does not block legacy routing, and a weekly allowance period by itself is not treated as evidence that the account is ineligible. +Administrators can override automatic media eligibility through the account create/update API by setting `extra.grok_media_eligible` to `false` (exclude) or `true` (force eligible). On update, set it to `null` to remove the override and return to automatic probe-based behavior; omitting the field preserves the current override. A weekly allowance period alone is not treated as a paid tier signal. Successful image responses must contain at least one actual image output; empty HTTP `200` responses trigger account failover instead of being counted and returned as successful generations. --- diff --git a/backend/cmd/cleanup-ingress-reject-logs/README.md b/backend/cmd/cleanup-ingress-reject-logs/README.md new file mode 100644 index 000000000..e2515a9dc --- /dev/null +++ b/backend/cmd/cleanup-ingress-reject-logs/README.md @@ -0,0 +1,20 @@ +# Ingress rejection log cleanup + +This maintenance command removes historical admission rejections from +`ops_error_logs` without matching unrelated authentication or upstream errors. +It is a dry run unless `--execute` is supplied, and always requires an explicit +RFC3339 cutoff. + +```sh +go run ./cmd/cleanup-ingress-reject-logs --before 2026-07-17T00:00:00Z +go run ./cmd/cleanup-ingress-reject-logs --before 2026-07-17T00:00:00Z --execute +``` + +Run the execute form only after every application instance has been upgraded so +older instances cannot add new ingress rejection rows below the chosen cutoff. +The classifier intentionally retains invariant failures such as +`USER_NOT_FOUND`, database errors, quota/billing errors, and upstream failures. + +After the rollout and cleanup are verified, run +`backend/scripts/finalize-ingress-reject-cleanup.sql` in a maintenance window to +remove the deprecated plaintext-key audit table and attribution columns. diff --git a/backend/cmd/cleanup-ingress-reject-logs/main.go b/backend/cmd/cleanup-ingress-reject-logs/main.go new file mode 100644 index 000000000..dbf045e0a --- /dev/null +++ b/backend/cmd/cleanup-ingress-reject-logs/main.go @@ -0,0 +1,218 @@ +package main + +import ( + "context" + "crypto/sha256" + "database/sql" + "encoding/hex" + "encoding/json" + "flag" + "fmt" + "log" + "sort" + "strings" + "time" + + _ "github.com/Wei-Shaw/sub2api/ent/runtime" + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/repository" + "github.com/lib/pq" +) + +const classifierVersion = "ingress-reject-v1" + +type candidate struct { + id int64 + statusCode int + message string + body string +} + +func main() { + beforeRaw := flag.String("before", "", "required RFC3339 cutoff; only older rows are considered") + execute := flag.Bool("execute", false, "delete matched rows (default is dry-run)") + batchSize := flag.Int("batch-size", 5000, "scan/delete batch size (1-5000)") + flag.Parse() + + if *beforeRaw == "" { + log.Fatal("--before is required") + } + before, err := time.Parse(time.RFC3339, *beforeRaw) + if err != nil { + log.Fatalf("invalid --before: %v", err) + } + if *batchSize < 1 || *batchSize > 5000 { + log.Fatal("--batch-size must be between 1 and 5000") + } + + cfg, err := config.LoadForBootstrap() + if err != nil { + log.Fatalf("load config: %v", err) + } + client, db, err := repository.InitEnt(cfg) + if err != nil { + log.Fatalf("initialize database: %v", err) + } + defer func() { _ = client.Close() }() + + ctx := context.Background() + counts, scanned, matched, deleted, err := cleanup(ctx, db, before, *batchSize, *execute) + if err != nil { + log.Fatalf("cleanup failed: %v", err) + } + + digest := sha256.Sum256([]byte(classifierVersion)) + mode := "dry-run" + if *execute { + mode = "execute" + } + fmt.Printf("mode=%s before=%s classifier=%s scanned=%d matched=%d deleted=%d\n", + mode, before.UTC().Format(time.RFC3339), hex.EncodeToString(digest[:]), scanned, matched, deleted) + reasons := make([]string, 0, len(counts)) + for reason := range counts { + reasons = append(reasons, reason) + } + sort.Strings(reasons) + for _, reason := range reasons { + fmt.Printf("reason=%s count=%d\n", reason, counts[reason]) + } + if *execute && deleted > 0 { + fmt.Println("cleanup complete; schedule VACUUM (ANALYZE) ops_error_logs during normal maintenance") + } +} + +func cleanup(ctx context.Context, db *sql.DB, before time.Time, batchSize int, execute bool) (map[string]int64, int64, int64, int64, error) { + counts := make(map[string]int64) + var cursor, scanned, matched, deleted int64 + for { + rows, err := db.QueryContext(ctx, ` + SELECT id, COALESCE(status_code, 0), COALESCE(error_message, ''), COALESCE(error_body, '') + FROM ops_error_logs + WHERE id > $1 + AND created_at < $2 + AND error_phase = 'auth' + AND account_id IS NULL + AND upstream_status_code IS NULL + AND COALESCE(upstream_error_message, '') = '' + AND COALESCE(upstream_error_detail, '') = '' + ORDER BY id ASC + LIMIT $3`, cursor, before, batchSize) + if err != nil { + return nil, scanned, matched, deleted, err + } + + batch := make([]candidate, 0, batchSize) + for rows.Next() { + var item candidate + if err := rows.Scan(&item.id, &item.statusCode, &item.message, &item.body); err != nil { + _ = rows.Close() + return nil, scanned, matched, deleted, err + } + batch = append(batch, item) + cursor = item.id + } + if err := rows.Err(); err != nil { + _ = rows.Close() + return nil, scanned, matched, deleted, err + } + _ = rows.Close() + if len(batch) == 0 { + break + } + + ids := make([]int64, 0, len(batch)) + for _, item := range batch { + scanned++ + if reason, ok := historicalIngressRejectReason(item); ok { + matched++ + counts[reason]++ + ids = append(ids, item.id) + } + } + if execute && len(ids) > 0 { + result, err := db.ExecContext(ctx, + `DELETE FROM ops_error_logs WHERE id = ANY($1) AND created_at < $2`, pq.Array(ids), before) + if err != nil { + return nil, scanned, matched, deleted, err + } + n, err := result.RowsAffected() + if err != nil { + return nil, scanned, matched, deleted, err + } + deleted += n + } + } + return counts, scanned, matched, deleted, nil +} + +func historicalIngressRejectReason(item candidate) (string, bool) { + code, message := parseErrorIdentity(item.body, item.message) + switch code { + case "API_KEY_REQUIRED": + return "missing_key", true + case "INVALID_API_KEY": + return "invalid_key", true + case "API_KEY_DISABLED": + return "key_disabled", true + case "USER_INACTIVE": + return "user_inactive", true + case "GROUP_DELETED": + return "group_deleted", true + case "GROUP_DISABLED": + return "group_disabled", true + case "GROUP_NOT_ALLOWED": + return "group_forbidden", true + case "ACCESS_DENIED": + return "ip_acl_denied", true + case "api_key_in_query_deprecated": + return "query_key_deprecated", true + } + + normalized := strings.TrimSpace(message) + switch { + case normalized == "API key is required": + return "missing_key", true + case normalized == "Invalid API key": + return "invalid_key", true + case normalized == "API key is disabled": + return "key_disabled", true + case normalized == "User account is not active": + return "user_inactive", true + case normalized == "API Key 所属分组已删除": + return "group_deleted", true + case normalized == "API Key 所属分组已停用": + return "group_disabled", true + case normalized == "API Key 所属专属分组不再允许当前用户使用": + return "group_forbidden", true + case normalized == "API Key is not assigned to any group and cannot be used. Please contact the administrator to assign it to a group.": + return "group_unassigned", true + case strings.HasPrefix(normalized, "Access denied. Your IP is "): + return "ip_acl_denied", true + case normalized == "Query parameter api_key is deprecated. Use Authorization header or key instead.": + return "query_key_deprecated", true + default: + return "", false + } +} + +func parseErrorIdentity(body, fallbackMessage string) (string, string) { + var payload struct { + Code string `json:"code"` + Message string `json:"message"` + Error struct { + Code json.RawMessage `json:"code"` + Message string `json:"message"` + } `json:"error"` + } + if err := json.Unmarshal([]byte(body), &payload); err != nil { + return "", fallbackMessage + } + message := payload.Message + if message == "" { + message = payload.Error.Message + } + if message == "" { + message = fallbackMessage + } + return strings.TrimSpace(payload.Code), message +} diff --git a/backend/cmd/cleanup-ingress-reject-logs/main_test.go b/backend/cmd/cleanup-ingress-reject-logs/main_test.go new file mode 100644 index 000000000..f7616ad2d --- /dev/null +++ b/backend/cmd/cleanup-ingress-reject-logs/main_test.go @@ -0,0 +1,29 @@ +package main + +import "testing" + +func TestHistoricalIngressRejectReason(t *testing.T) { + tests := []struct { + name string + item candidate + reason string + match bool + }{ + {name: "standard invalid key", item: candidate{body: `{"code":"INVALID_API_KEY","message":"Invalid API key"}`}, reason: "invalid_key", match: true}, + {name: "google missing key", item: candidate{body: `{"error":{"code":401,"message":"API key is required","status":"UNAUTHENTICATED"}}`}, reason: "missing_key", match: true}, + {name: "google group deleted", item: candidate{body: `{"error":{"code":403,"message":"API Key 所属分组已删除","status":"PERMISSION_DENIED"}}`}, reason: "group_deleted", match: true}, + {name: "ip acl", item: candidate{body: `{"code":"ACCESS_DENIED","message":"Access denied. Your IP is 192.0.2.1"}`}, reason: "ip_acl_denied", match: true}, + {name: "user not found remains", item: candidate{body: `{"code":"USER_NOT_FOUND","message":"User associated with API key not found"}`}, match: false}, + {name: "quota remains", item: candidate{body: `{"code":"API_KEY_QUOTA_EXHAUSTED","message":"quota"}`}, match: false}, + {name: "database failure remains", item: candidate{statusCode: 500, message: "Failed to validate API key", body: `{"code":"INTERNAL_ERROR","message":"Failed to validate API key"}`}, match: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + reason, ok := historicalIngressRejectReason(tt.item) + if ok != tt.match || reason != tt.reason { + t.Fatalf("got (%q, %v), want (%q, %v)", reason, ok, tt.reason, tt.match) + } + }) + } +} diff --git a/backend/cmd/server/wire.go b/backend/cmd/server/wire.go index cd4454f52..22d7322af 100644 --- a/backend/cmd/server/wire.go +++ b/backend/cmd/server/wire.go @@ -81,6 +81,10 @@ func provideCleanup( opsCleanup *service.OpsCleanupService, opsScheduledReport *service.OpsScheduledReportService, opsSystemLogSink *service.OpsSystemLogSink, + opsService *service.OpsService, + opsIngressReject *service.OpsIngressRejectAggregator, + apiKeyService *service.APIKeyService, + authCacheInvalidationWorker *service.AuthCacheInvalidationWorker, schedulerSnapshot *service.SchedulerSnapshotService, tokenRefresh *service.TokenRefreshService, accountExpiry *service.AccountExpiryService, @@ -121,6 +125,30 @@ func provideCleanup( // 应用层清理步骤可并行执行,基础设施资源(Redis/Ent)最后按顺序关闭。 parallelSteps := []cleanupStep{ + {"OpsIngressRejectAggregator", func() error { + if opsIngressReject != nil { + opsIngressReject.Stop() + } + return nil + }}, + {"AuthCacheInvalidationWorker", func() error { + if authCacheInvalidationWorker != nil { + authCacheInvalidationWorker.Stop() + } + return nil + }}, + {"AuthCacheInvalidationSubscriber", func() error { + if apiKeyService != nil { + apiKeyService.StopAuthCacheInvalidationSubscriber() + } + return nil + }}, + {"OpsRuntimeSettingsRefresh", func() error { + if opsService != nil { + opsService.StopRuntimeSettingsRefresh() + } + return nil + }}, {"PromptAuditService", func() error { if promptAudit != nil { return promptAudit.Shutdown(ctx) diff --git a/backend/cmd/server/wire_gen.go b/backend/cmd/server/wire_gen.go index 7a88fa355..39354c1cf 100644 --- a/backend/cmd/server/wire_gen.go +++ b/backend/cmd/server/wire_gen.go @@ -157,7 +157,9 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { antigravityGatewayService := service.NewAntigravityGatewayService(accountRepository, gatewayCache, schedulerSnapshotService, antigravityTokenProvider, rateLimitService, httpUpstream, settingService, internal500CounterCache) geminiMessagesCompatService := service.NewGeminiMessagesCompatService(accountRepository, groupRepository, gatewayCache, schedulerSnapshotService, geminiTokenProvider, rateLimitService, httpUpstream, antigravityGatewayService, configConfig) opsSystemLogSink := service.ProvideOpsSystemLogSink(opsRepository) - opsService := service.ProvideOpsService(opsRepository, settingRepository, configConfig, accountRepository, userRepository, concurrencyService, gatewayService, openAIGatewayService, geminiMessagesCompatService, antigravityGatewayService, opsSystemLogSink, settingService) + authCacheInvalidationOutboxRepository := repository.NewAuthCacheInvalidationOutboxRepository(db) + authCacheInvalidationWorker := service.ProvideAuthCacheInvalidationWorker(authCacheInvalidationOutboxRepository, apiKeyCache, apiKeyService) + opsService := service.ProvideOpsService(opsRepository, settingRepository, configConfig, accountRepository, userRepository, concurrencyService, gatewayService, openAIGatewayService, geminiMessagesCompatService, antigravityGatewayService, opsSystemLogSink, settingService, authCacheInvalidationWorker, apiKeyService) usageHandler := handler.NewUsageHandler(usageService, apiKeyService, opsService, settingService) redeemHandler := handler.NewRedeemHandler(redeemService) subscriptionHandler := handler.NewSubscriptionHandler(subscriptionService) @@ -270,7 +272,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { legacyEngine := securityaudit.NewLegacyModerationAdapter(contentModerationService) coordinator := securityaudit.NewCoordinator(legacyEngine, promptService) gatewayHandler := handler.ProvideGatewayHandler(gatewayService, openAIGatewayService, geminiMessagesCompatService, antigravityGatewayService, userService, concurrencyService, billingCacheService, usageService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, userMessageQueueService, configConfig, settingService, coordinator) - openAIGatewayHandler := handler.ProvideOpenAIGatewayHandler(openAIGatewayService, concurrencyService, billingCacheService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, opsService, configConfig, coordinator) + openAIGatewayHandler := handler.ProvideOpenAIGatewayHandler(openAIGatewayService, concurrencyService, billingCacheService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, opsService, grokQuotaService, configConfig, coordinator) handlerSettingHandler := handler.ProvideSettingHandler(settingService, buildInfo, notificationEmailService) totpHandler := handler.NewTotpHandler(totpService) handlerPaymentHandler := handler.NewPaymentHandler(paymentService, paymentConfigService) @@ -306,6 +308,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { opsAlertEvaluatorService := service.ProvideOpsAlertEvaluatorService(opsService, opsRepository, emailService, redisClient, configConfig, proxyRepository) opsCleanupService := service.ProvideOpsCleanupService(opsRepository, db, redisClient, configConfig, channelMonitorService, settingRepository, opsService) opsScheduledReportService := service.ProvideOpsScheduledReportService(opsService, userService, emailService, redisClient, configConfig) + opsIngressRejectAggregator := service.ProvideOpsIngressRejectAggregator(opsRepository, opsService) accountExpiryService := service.ProvideAccountExpiryService(accountRepository) proxyExpiryService := service.ProvideProxyExpiryService(proxyRepository) subscriptionExpiryService := service.ProvideSubscriptionExpiryService(userSubscriptionRepository, settingRepository, notificationEmailService, leaderLockCache, db) @@ -314,7 +317,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { paymentOrderExpiryService := service.ProvidePaymentOrderExpiryService(paymentService, leaderLockCache, db) channelMonitorRunner := service.ProvideChannelMonitorRunner(channelMonitorService, settingService) userPlatformQuotaUsageFlusher := service.ProvideUserPlatformQuotaUsageFlusher(configConfig, billingCache, serviceUserPlatformQuotaRepository, timingWheelService) - v := provideCleanup(client, redisClient, opsMetricsCollector, opsAggregationService, opsAlertEvaluatorService, opsCleanupService, opsScheduledReportService, opsSystemLogSink, schedulerSnapshotService, tokenRefreshService, accountExpiryService, proxyExpiryService, subscriptionExpiryService, usageCleanupService, idempotencyCleanupService, batchImageCleanupService, batchImageWorkerRuntime, pricingService, emailQueueService, billingCacheService, usageRecordWorkerPool, subscriptionService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, grokOAuthService, openAIGatewayService, scheduledTestRunnerService, backupService, paymentOrderExpiryService, channelMonitorRunner, userPlatformQuotaUsageFlusher, upstreamBillingProbeService, auditLogService, promptService) + v := provideCleanup(client, redisClient, opsMetricsCollector, opsAggregationService, opsAlertEvaluatorService, opsCleanupService, opsScheduledReportService, opsSystemLogSink, opsService, opsIngressRejectAggregator, apiKeyService, authCacheInvalidationWorker, schedulerSnapshotService, tokenRefreshService, accountExpiryService, proxyExpiryService, subscriptionExpiryService, usageCleanupService, idempotencyCleanupService, batchImageCleanupService, batchImageWorkerRuntime, pricingService, emailQueueService, billingCacheService, usageRecordWorkerPool, subscriptionService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, grokOAuthService, openAIGatewayService, scheduledTestRunnerService, backupService, paymentOrderExpiryService, channelMonitorRunner, userPlatformQuotaUsageFlusher, upstreamBillingProbeService, auditLogService, promptService) application := &Application{ Server: httpServer, PromptAudit: promptService, @@ -351,6 +354,10 @@ func provideCleanup( opsCleanup *service.OpsCleanupService, opsScheduledReport *service.OpsScheduledReportService, opsSystemLogSink *service.OpsSystemLogSink, + opsService *service.OpsService, + opsIngressReject *service.OpsIngressRejectAggregator, + apiKeyService *service.APIKeyService, + authCacheInvalidationWorker *service.AuthCacheInvalidationWorker, schedulerSnapshot *service.SchedulerSnapshotService, tokenRefresh *service.TokenRefreshService, accountExpiry *service.AccountExpiryService, @@ -390,6 +397,30 @@ func provideCleanup( } parallelSteps := []cleanupStep{ + {"OpsIngressRejectAggregator", func() error { + if opsIngressReject != nil { + opsIngressReject.Stop() + } + return nil + }}, + {"AuthCacheInvalidationWorker", func() error { + if authCacheInvalidationWorker != nil { + authCacheInvalidationWorker.Stop() + } + return nil + }}, + {"AuthCacheInvalidationSubscriber", func() error { + if apiKeyService != nil { + apiKeyService.StopAuthCacheInvalidationSubscriber() + } + return nil + }}, + {"OpsRuntimeSettingsRefresh", func() error { + if opsService != nil { + opsService.StopRuntimeSettingsRefresh() + } + return nil + }}, {"PromptAuditService", func() error { if promptAudit != nil { return promptAudit.Shutdown(ctx) diff --git a/backend/cmd/server/wire_gen_test.go b/backend/cmd/server/wire_gen_test.go index f9a479bc6..3f159af7f 100644 --- a/backend/cmd/server/wire_gen_test.go +++ b/backend/cmd/server/wire_gen_test.go @@ -58,6 +58,10 @@ func TestProvideCleanup_WithMinimalDependencies_NoPanic(t *testing.T) { &service.OpsCleanupService{}, &service.OpsScheduledReportService{}, opsSystemLogSinkSvc, + nil, // opsService + nil, // opsIngressRejectAggregator + nil, // apiKeyService + nil, // authCacheInvalidationWorker schedulerSnapshotSvc, tokenRefreshSvc, accountExpirySvc, diff --git a/backend/internal/config/config.go b/backend/internal/config/config.go index def68f3a2..38ebc9427 100644 --- a/backend/internal/config/config.go +++ b/backend/internal/config/config.go @@ -645,6 +645,7 @@ type ServerConfig struct { EnableServerTiming bool `mapstructure:"enable_server_timing"` // Admin UI Server-Timing response header FrontendURL string `mapstructure:"frontend_url"` // 前端基础 URL,用于生成邮件中的外部链接 ReadHeaderTimeout int `mapstructure:"read_header_timeout"` // 读取请求头超时(秒) + MaxHeaderBytes int `mapstructure:"max_header_bytes"` // 请求头最大字节数(HTTP/2 映射为 header-list 上限) IdleTimeout int `mapstructure:"idle_timeout"` // 空闲连接超时(秒) TrustedProxies []string `mapstructure:"trusted_proxies"` // 可信代理列表(CIDR/IP) MaxRequestBodySize int64 `mapstructure:"max_request_body_size"` // 全局最大请求体限制 @@ -796,6 +797,8 @@ type GatewayConfig struct { OpenAIHighEffortFirstOutputTimeoutSeconds int `mapstructure:"openai_high_effort_first_output_timeout_seconds"` // 请求体最大字节数,用于网关请求体大小限制 MaxBodySize int64 `mapstructure:"max_body_size"` + // TextMaxBodySize limits endpoints that cannot carry inline image/video payloads. + TextMaxBodySize int64 `mapstructure:"text_max_body_size"` // 非流式上游响应体读取上限(字节),用于防止无界读取导致内存放大 UpstreamResponseReadMaxBytes int64 `mapstructure:"upstream_response_read_max_bytes"` // 代理探测响应体读取上限(字节) @@ -1419,12 +1422,22 @@ type RateLimitConfig struct { // APIKeyAuthCacheConfig API Key 认证缓存配置 type APIKeyAuthCacheConfig struct { - L1Size int `mapstructure:"l1_size"` - L1TTLSeconds int `mapstructure:"l1_ttl_seconds"` - L2TTLSeconds int `mapstructure:"l2_ttl_seconds"` - NegativeTTLSeconds int `mapstructure:"negative_ttl_seconds"` - JitterPercent int `mapstructure:"jitter_percent"` - Singleflight bool `mapstructure:"singleflight"` + L1Size int `mapstructure:"l1_size"` + L1TTLSeconds int `mapstructure:"l1_ttl_seconds"` + L2TTLSeconds int `mapstructure:"l2_ttl_seconds"` + NegativeTTLSeconds int `mapstructure:"negative_ttl_seconds"` + JitterPercent int `mapstructure:"jitter_percent"` + Singleflight bool `mapstructure:"singleflight"` + LookupConcurrency int `mapstructure:"lookup_concurrency"` + InvalidAbuse InvalidAuthAbuseConfig `mapstructure:"invalid_abuse"` +} + +type InvalidAuthAbuseConfig struct { + Enabled bool `mapstructure:"enabled"` + Threshold int `mapstructure:"threshold"` + WindowSeconds int `mapstructure:"window_seconds"` + BlockSeconds int `mapstructure:"block_seconds"` + Capacity int `mapstructure:"capacity"` } // SubscriptionCacheConfig 订阅认证 L1 缓存配置 @@ -1698,8 +1711,9 @@ func setDefaults() { viper.SetDefault("server.mode", "release") viper.SetDefault("server.enable_server_timing", false) viper.SetDefault("server.frontend_url", "") - viper.SetDefault("server.read_header_timeout", 30) // 30秒读取请求头 - viper.SetDefault("server.idle_timeout", 120) // 120秒空闲超时 + viper.SetDefault("server.read_header_timeout", 10) // 10秒读取请求头 + viper.SetDefault("server.max_header_bytes", 64*1024) + viper.SetDefault("server.idle_timeout", 120) // 120秒空闲超时 viper.SetDefault("server.trusted_proxies", []string{}) viper.SetDefault("server.max_request_body_size", int64(256*1024*1024)) // H2C 默认配置 @@ -1983,6 +1997,12 @@ func setDefaults() { viper.SetDefault("api_key_auth_cache.negative_ttl_seconds", 30) viper.SetDefault("api_key_auth_cache.jitter_percent", 10) viper.SetDefault("api_key_auth_cache.singleflight", true) + viper.SetDefault("api_key_auth_cache.lookup_concurrency", 64) + viper.SetDefault("api_key_auth_cache.invalid_abuse.enabled", true) + viper.SetDefault("api_key_auth_cache.invalid_abuse.threshold", 120) + viper.SetDefault("api_key_auth_cache.invalid_abuse.window_seconds", 60) + viper.SetDefault("api_key_auth_cache.invalid_abuse.block_seconds", 60) + viper.SetDefault("api_key_auth_cache.invalid_abuse.capacity", 16384) // Subscription auth L1 cache viper.SetDefault("subscription_cache.l1_size", 16384) @@ -2111,6 +2131,7 @@ func setDefaults() { viper.SetDefault("gateway.antigravity_fallback_cooldown_minutes", 1) viper.SetDefault("gateway.antigravity_extra_retries", 10) viper.SetDefault("gateway.max_body_size", int64(256*1024*1024)) + viper.SetDefault("gateway.text_max_body_size", int64(32*1024*1024)) viper.SetDefault("gateway.upstream_response_read_max_bytes", DefaultUpstreamResponseReadMaxBytes) viper.SetDefault("gateway.proxy_probe_response_read_max_bytes", int64(1024*1024)) viper.SetDefault("gateway.gemini_debug_response_headers", false) @@ -2208,6 +2229,49 @@ func setDefaults() { } func (c *Config) Validate() error { + if c.Server.ReadHeaderTimeout < 1 || c.Server.ReadHeaderTimeout > 60 { + return fmt.Errorf("server.read_header_timeout must be between 1 and 60 seconds") + } + if c.Server.MaxHeaderBytes < 8*1024 || c.Server.MaxHeaderBytes > 1024*1024 { + return fmt.Errorf("server.max_header_bytes must be between 8192 and 1048576 bytes") + } + if c.Server.IdleTimeout <= 0 { + return fmt.Errorf("server.idle_timeout must be positive") + } + if c.Server.MaxRequestBodySize < 0 { + return fmt.Errorf("server.max_request_body_size must be non-negative") + } + if c.Server.H2C.Enabled { + if c.Server.H2C.MaxConcurrentStreams == 0 { + return fmt.Errorf("server.h2c.max_concurrent_streams must be positive") + } + if c.Server.H2C.IdleTimeout <= 0 { + return fmt.Errorf("server.h2c.idle_timeout must be positive") + } + if c.Server.H2C.MaxReadFrameSize < 16*1024 || c.Server.H2C.MaxReadFrameSize > 16*1024*1024-1 { + return fmt.Errorf("server.h2c.max_read_frame_size must be between 16384 and 16777215 bytes") + } + if c.Server.H2C.MaxUploadBufferPerConnection < 65535 { + return fmt.Errorf("server.h2c.max_upload_buffer_per_connection must be at least 65535 bytes") + } + if c.Server.H2C.MaxUploadBufferPerStream <= 0 { + return fmt.Errorf("server.h2c.max_upload_buffer_per_stream must be positive") + } + } + if c.APIKeyAuth.InvalidAbuse.Enabled { + if c.APIKeyAuth.InvalidAbuse.Threshold < 10 { + return fmt.Errorf("api_key_auth_cache.invalid_abuse.threshold must be at least 10") + } + if c.APIKeyAuth.InvalidAbuse.WindowSeconds < 1 || c.APIKeyAuth.InvalidAbuse.WindowSeconds > 3600 { + return fmt.Errorf("api_key_auth_cache.invalid_abuse.window_seconds must be between 1 and 3600") + } + if c.APIKeyAuth.InvalidAbuse.BlockSeconds < 1 || c.APIKeyAuth.InvalidAbuse.BlockSeconds > 3600 { + return fmt.Errorf("api_key_auth_cache.invalid_abuse.block_seconds must be between 1 and 3600") + } + if c.APIKeyAuth.InvalidAbuse.Capacity < 256 || c.APIKeyAuth.InvalidAbuse.Capacity > 1_000_000 { + return fmt.Errorf("api_key_auth_cache.invalid_abuse.capacity must be between 256 and 1000000") + } + } jwtSecret := strings.TrimSpace(c.JWT.Secret) if jwtSecret == "" { return fmt.Errorf("jwt.secret is required") @@ -2733,6 +2797,9 @@ func (c *Config) Validate() error { if c.Gateway.MaxBodySize <= 0 { return fmt.Errorf("gateway.max_body_size must be positive") } + if c.Gateway.TextMaxBodySize <= 0 || c.Gateway.TextMaxBodySize > c.Gateway.MaxBodySize { + return fmt.Errorf("gateway.text_max_body_size must be positive and no greater than gateway.max_body_size") + } if c.Gateway.UpstreamResponseReadMaxBytes <= 0 { return fmt.Errorf("gateway.upstream_response_read_max_bytes must be positive") } diff --git a/backend/internal/config/config_test.go b/backend/internal/config/config_test.go index badfcf6ac..36ae89ffa 100644 --- a/backend/internal/config/config_test.go +++ b/backend/internal/config/config_test.go @@ -35,6 +35,18 @@ func TestLoadServerTimingConfig(t *testing.T) { }) } +func TestLoadHTTPIngressSafetyDefaults(t *testing.T) { + resetViperWithJWTSecret(t) + cfg, err := Load() + require.NoError(t, err) + require.Equal(t, 10, cfg.Server.ReadHeaderTimeout) + require.Equal(t, 64*1024, cfg.Server.MaxHeaderBytes) + require.Equal(t, int64(32*1024*1024), cfg.Gateway.TextMaxBodySize) + require.True(t, cfg.APIKeyAuth.InvalidAbuse.Enabled) + require.Equal(t, 120, cfg.APIKeyAuth.InvalidAbuse.Threshold) + require.Equal(t, 16384, cfg.APIKeyAuth.InvalidAbuse.Capacity) +} + func TestLoadForBootstrapAllowsMissingJWTSecret(t *testing.T) { viper.Reset() t.Setenv("JWT_SECRET", "") @@ -1185,6 +1197,51 @@ func TestValidateConfigErrors(t *testing.T) { mutate func(*Config) wantErr string }{ + { + name: "server read header timeout", + mutate: func(c *Config) { c.Server.ReadHeaderTimeout = 0 }, + wantErr: "server.read_header_timeout", + }, + { + name: "server max header bytes too small", + mutate: func(c *Config) { c.Server.MaxHeaderBytes = 4096 }, + wantErr: "server.max_header_bytes", + }, + { + name: "server max request body size", + mutate: func(c *Config) { c.Server.MaxRequestBodySize = -1 }, + wantErr: "server.max_request_body_size", + }, + { + name: "h2c zero concurrent streams", + mutate: func(c *Config) { + c.Server.H2C.Enabled = true + c.Server.H2C.MaxConcurrentStreams = 0 + }, + wantErr: "server.h2c.max_concurrent_streams", + }, + { + name: "h2c oversized read frame", + mutate: func(c *Config) { + c.Server.H2C.Enabled = true + c.Server.H2C.MaxReadFrameSize = 16 * 1024 * 1024 + }, + wantErr: "server.h2c.max_read_frame_size", + }, + { + name: "invalid auth abuse threshold too small", + mutate: func(c *Config) { + c.APIKeyAuth.InvalidAbuse.Threshold = 9 + }, + wantErr: "api_key_auth_cache.invalid_abuse.threshold", + }, + { + name: "invalid auth abuse capacity too small", + mutate: func(c *Config) { + c.APIKeyAuth.InvalidAbuse.Capacity = 255 + }, + wantErr: "api_key_auth_cache.invalid_abuse.capacity", + }, { name: "jwt secret required", mutate: func(c *Config) { c.JWT.Secret = "" }, @@ -1386,6 +1443,11 @@ func TestValidateConfigErrors(t *testing.T) { mutate: func(c *Config) { c.Gateway.MaxBodySize = 0 }, wantErr: "gateway.max_body_size", }, + { + name: "gateway text body exceeds media body", + mutate: func(c *Config) { c.Gateway.TextMaxBodySize = c.Gateway.MaxBodySize + 1 }, + wantErr: "gateway.text_max_body_size", + }, { name: "gateway response header timeout", mutate: func(c *Config) { c.Gateway.ResponseHeaderTimeout = -1 }, diff --git a/backend/internal/handler/admin/account_handler.go b/backend/internal/handler/admin/account_handler.go index d9b742003..6f585fcfc 100644 --- a/backend/internal/handler/admin/account_handler.go +++ b/backend/internal/handler/admin/account_handler.go @@ -61,7 +61,7 @@ type AccountHandler struct { sessionLimitCache service.SessionLimitCache rpmCache service.RPMCache tokenCacheInvalidator service.TokenCacheInvalidator - grokImportProber grokUsageProber + grokImportProber grokImportProber upstreamBillingProbe *service.UpstreamBillingProbeService } @@ -121,6 +121,7 @@ type CreateAccountRequest struct { GroupIDs []int64 `json:"group_ids"` ExpiresAt *int64 `json:"expires_at"` AutoPauseOnExpired *bool `json:"auto_pause_on_expired"` + ProbeEnabled *bool `json:"upstream_billing_probe_enabled"` ConfirmMixedChannelRisk *bool `json:"confirm_mixed_channel_risk"` // 用户确认混合渠道风险 } @@ -159,6 +160,7 @@ type BulkUpdateAccountsRequest struct { GroupIDs *[]int64 `json:"group_ids"` Credentials map[string]any `json:"credentials"` Extra map[string]any `json:"extra"` + ProbeEnabled *bool `json:"upstream_billing_probe_enabled"` ConfirmMixedChannelRisk *bool `json:"confirm_mixed_channel_risk"` // 用户确认混合渠道风险 } @@ -828,6 +830,7 @@ func (h *AccountHandler) Create(c *gin.Context) { GroupIDs: req.GroupIDs, ExpiresAt: req.ExpiresAt, AutoPauseOnExpired: req.AutoPauseOnExpired, + ProbeEnabled: req.ProbeEnabled, SkipMixedChannelCheck: skipCheck, }) if execErr != nil { @@ -1907,7 +1910,8 @@ func (h *AccountHandler) BulkUpdate(c *gin.Context) { req.Schedulable != nil || req.GroupIDs != nil || len(req.Credentials) > 0 || - len(req.Extra) > 0 + len(req.Extra) > 0 || + req.ProbeEnabled != nil if !hasUpdates { response.BadRequest(c, "No updates provided") @@ -1928,6 +1932,7 @@ func (h *AccountHandler) BulkUpdate(c *gin.Context) { GroupIDs: req.GroupIDs, Credentials: req.Credentials, Extra: req.Extra, + ProbeEnabled: req.ProbeEnabled, SkipMixedChannelCheck: skipCheck, }) if err != nil { diff --git a/backend/internal/handler/admin/account_handler_mixed_channel_test.go b/backend/internal/handler/admin/account_handler_mixed_channel_test.go index a1dc25124..bc4a3cf49 100644 --- a/backend/internal/handler/admin/account_handler_mixed_channel_test.go +++ b/backend/internal/handler/admin/account_handler_mixed_channel_test.go @@ -222,3 +222,22 @@ func TestBulkUpdateAcceptsFilterTargetRequest(t *testing.T) { require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp)) require.Equal(t, float64(0), resp["code"]) } + +func TestBulkUpdateAcceptsDedicatedUpstreamBillingProbeSetting(t *testing.T) { + adminSvc := newStubAdminService() + router := setupAccountMixedChannelRouter(adminSvc) + + body, _ := json.Marshal(map[string]any{ + "account_ids": []int64{1, 2}, + "upstream_billing_probe_enabled": false, + }) + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/bulk-update", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + router.ServeHTTP(rec, req) + + require.Equal(t, http.StatusOK, rec.Code) + require.NotNil(t, adminSvc.lastBulkUpdateAccountInput) + require.NotNil(t, adminSvc.lastBulkUpdateAccountInput.ProbeEnabled) + require.False(t, *adminSvc.lastBulkUpdateAccountInput.ProbeEnabled) +} diff --git a/backend/internal/handler/admin/admin_service_stub_test.go b/backend/internal/handler/admin/admin_service_stub_test.go index f5c238a61..035e171c9 100644 --- a/backend/internal/handler/admin/admin_service_stub_test.go +++ b/backend/internal/handler/admin/admin_service_stub_test.go @@ -33,6 +33,7 @@ type stubAdminService struct { createSparkShadowErr error updateAccountErr error bulkUpdateAccountErr error + lastBulkUpdateAccountInput *service.BulkUpdateAccountsInput getAccountResult *service.Account updateAccountCalls int updateAccountExtraCalls int @@ -478,6 +479,7 @@ func (s *stubAdminService) SetAccountSchedulable(ctx context.Context, id int64, } func (s *stubAdminService) BulkUpdateAccounts(ctx context.Context, input *service.BulkUpdateAccountsInput) (*service.BulkUpdateAccountsResult, error) { + s.lastBulkUpdateAccountInput = input if s.bulkUpdateAccountErr != nil { return nil, s.bulkUpdateAccountErr } diff --git a/backend/internal/handler/admin/grok_import_probe.go b/backend/internal/handler/admin/grok_import_probe.go index b42957ce6..9ef4ce8cd 100644 --- a/backend/internal/handler/admin/grok_import_probe.go +++ b/backend/internal/handler/admin/grok_import_probe.go @@ -15,12 +15,12 @@ const ( grokImportProbeTimeout = 25 * time.Second ) -type grokUsageProber interface { - ProbeUsage(ctx context.Context, accountID int64) (*service.GrokQuotaProbeResult, error) +type grokImportProber interface { + QueryQuota(ctx context.Context, accountID int64) (*service.GrokQuotaProbeResult, error) } type grokImportProbeTask struct { - prober grokUsageProber + prober grokImportProber accountID int64 } @@ -51,7 +51,7 @@ func newGrokImportProbeScheduler(concurrency int, timeout time.Duration) *grokIm } } -func (s *grokImportProbeScheduler) schedule(prober grokUsageProber, account *service.Account) { +func (s *grokImportProbeScheduler) schedule(prober grokImportProber, account *service.Account) { if s == nil || prober == nil || account == nil || account.ID <= 0 { return } @@ -97,7 +97,7 @@ func (s *grokImportProbeScheduler) nextTask() (grokImportProbeTask, bool) { return task, true } -func (s *grokImportProbeScheduler) run(prober grokUsageProber, accountID int64) { +func (s *grokImportProbeScheduler) run(prober grokImportProber, accountID int64) { defer func() { if recovered := recover(); recovered != nil { slog.Error( @@ -112,7 +112,7 @@ func (s *grokImportProbeScheduler) run(prober grokUsageProber, accountID int64) // while this timeout only bounds the actual upstream probe execution. ctx, cancel := context.WithTimeout(context.Background(), s.timeout) defer cancel() - result, err := prober.ProbeUsage(ctx, accountID) + result, err := prober.QueryQuota(ctx, accountID) if err != nil { slog.Warn( "grok_import_active_probe_failed", diff --git a/backend/internal/handler/admin/grok_import_probe_test.go b/backend/internal/handler/admin/grok_import_probe_test.go index 3b8fc0ca6..255f04d0d 100644 --- a/backend/internal/handler/admin/grok_import_probe_test.go +++ b/backend/internal/handler/admin/grok_import_probe_test.go @@ -36,7 +36,7 @@ func newGrokImportProbeStub(buffer int) *grokImportProbeStub { } } -func (s *grokImportProbeStub) ProbeUsage(ctx context.Context, accountID int64) (*service.GrokQuotaProbeResult, error) { +func (s *grokImportProbeStub) QueryQuota(ctx context.Context, accountID int64) (*service.GrokQuotaProbeResult, error) { _, deadlineSeen := ctx.Deadline() s.mu.Lock() s.calls[accountID]++ @@ -69,7 +69,7 @@ func (s *grokImportProbeStub) ProbeUsage(ctx context.Context, accountID int64) ( return nil, failure } return &service.GrokQuotaProbeResult{ - Source: "active_probe", + Source: "hybrid_probe", Model: "grok-4.5", StatusCode: 200, ResetSupported: false, diff --git a/backend/internal/handler/admin/grok_oauth_handler.go b/backend/internal/handler/admin/grok_oauth_handler.go index ac05fd8a6..1f679b956 100644 --- a/backend/internal/handler/admin/grok_oauth_handler.go +++ b/backend/internal/handler/admin/grok_oauth_handler.go @@ -23,7 +23,7 @@ type GrokOAuthHandler struct { grokOAuthService *service.GrokOAuthService adminService service.AdminService quotaService *service.GrokQuotaService - importProber grokUsageProber + importProber grokImportProber reconciler service.GrokOAuthReconciler } diff --git a/backend/internal/handler/admin/ops_auth_cache_health_handler.go b/backend/internal/handler/admin/ops_auth_cache_health_handler.go new file mode 100644 index 000000000..316707373 --- /dev/null +++ b/backend/internal/handler/admin/ops_auth_cache_health_handler.go @@ -0,0 +1,21 @@ +package admin + +import ( + "net/http" + + "github.com/Wei-Shaw/sub2api/internal/pkg/response" + "github.com/gin-gonic/gin" +) + +// GetAuthCacheInvalidationHealth exposes durable outbox lag and subscriber health. +func (h *OpsHandler) GetAuthCacheInvalidationHealth(c *gin.Context) { + if h.opsService == nil { + response.Error(c, http.StatusServiceUnavailable, "Ops service not available") + return + } + if err := h.opsService.RequireMonitoringEnabled(c.Request.Context()); err != nil { + response.ErrorFrom(c, err) + return + } + response.Success(c, h.opsService.GetAuthCacheInvalidationHealth(c.Request.Context())) +} diff --git a/backend/internal/handler/admin/ops_ingress_reject_handler.go b/backend/internal/handler/admin/ops_ingress_reject_handler.go new file mode 100644 index 000000000..6b45a7d99 --- /dev/null +++ b/backend/internal/handler/admin/ops_ingress_reject_handler.go @@ -0,0 +1,131 @@ +package admin + +import ( + "net/http" + "net/netip" + "strconv" + "strings" + + "github.com/Wei-Shaw/sub2api/internal/pkg/response" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" +) + +var ingressRejectReasons = map[string]struct{}{ + "query_api_key_deprecated": {}, "api_key_required": {}, "invalid_api_key": {}, + "invalid_auth_rate_limited": {}, + "api_key_auth_overloaded": {}, + "api_key_disabled": {}, "ip_restricted": {}, "user_inactive": {}, "group_deleted": {}, + "group_disabled": {}, "group_not_allowed": {}, "group_unassigned": {}, "other": {}, +} + +var ingressRejectRouteFamilies = map[string]struct{}{ + "antigravity": {}, "gemini": {}, "codex": {}, "messages": {}, "responses": {}, + "chat_completions": {}, "images": {}, "videos": {}, "embeddings": {}, "models": {}, "other": {}, +} + +var ingressRejectProtocols = map[string]struct{}{ + "google": {}, "anthropic": {}, "openai": {}, "gateway": {}, "other": {}, +} + +// ListIngressRejects returns bounded security aggregates, never raw credentials or request bodies. +func (h *OpsHandler) ListIngressRejects(c *gin.Context) { + if h.opsService == nil { + response.Error(c, http.StatusServiceUnavailable, "Ops service not available") + return + } + page, pageSize := response.ParsePagination(c) + if pageSize > 200 { + pageSize = 200 + } + startTime, endTime, err := parseOpsTimeRange(c, "1h") + if err != nil { + response.BadRequest(c, err.Error()) + return + } + filter := &service.OpsIngressRejectFilter{Page: page, PageSize: pageSize} + if !startTime.IsZero() { + filter.StartTime = &startTime + } + if !endTime.IsZero() { + filter.EndTime = &endTime + } + if filter.RejectReason, err = parseIngressRejectEnum(c, "reason", ingressRejectReasons); err != nil { + response.BadRequest(c, err.Error()) + return + } + if filter.RouteFamily, err = parseIngressRejectEnum(c, "route_family", ingressRejectRouteFamilies); err != nil { + response.BadRequest(c, err.Error()) + return + } + if filter.Protocol, err = parseIngressRejectEnum(c, "protocol", ingressRejectProtocols); err != nil { + response.BadRequest(c, err.Error()) + return + } + if raw := strings.TrimSpace(c.Query("client_ip")); raw != "" { + addr, parseErr := netip.ParseAddr(raw) + if parseErr != nil { + response.BadRequest(c, "Invalid client_ip") + return + } + addr = addr.Unmap() + if addr.Is6() { + addr = netip.PrefixFrom(addr, 64).Masked().Addr() + } + filter.ClientIP = addr.String() + } + if filter.UserID, err = parseOptionalPositiveID(c, "user_id"); err != nil { + response.BadRequest(c, err.Error()) + return + } + if filter.APIKeyID, err = parseOptionalPositiveID(c, "api_key_id"); err != nil { + response.BadRequest(c, err.Error()) + return + } + + result, err := h.opsService.ListIngressRejects(c.Request.Context(), filter) + if err != nil { + response.ErrorFrom(c, err) + return + } + response.Success(c, result) +} + +func (h *OpsHandler) GetIngressRejectHealth(c *gin.Context) { + if h.opsService == nil { + response.Error(c, http.StatusServiceUnavailable, "Ops service not available") + return + } + if err := h.opsService.RequireMonitoringEnabled(c.Request.Context()); err != nil { + response.ErrorFrom(c, err) + return + } + response.Success(c, h.opsService.GetIngressRejectHealth()) +} + +func parseIngressRejectEnum(c *gin.Context, name string, allowed map[string]struct{}) (string, error) { + value := strings.TrimSpace(c.Query(name)) + if value == "" { + return "", nil + } + if _, ok := allowed[value]; !ok { + return "", &ingressRejectQueryError{message: "Invalid " + name} + } + return value, nil +} + +func parseOptionalPositiveID(c *gin.Context, name string) (*int64, error) { + raw := strings.TrimSpace(c.Query(name)) + if raw == "" { + return nil, nil + } + value, err := strconv.ParseInt(raw, 10, 64) + if err != nil || value <= 0 { + return nil, &ingressRejectQueryError{message: "Invalid " + name} + } + return &value, nil +} + +type ingressRejectQueryError struct{ message string } + +func (e *ingressRejectQueryError) Error() string { return e.message } diff --git a/backend/internal/handler/admin/ops_ingress_reject_handler_test.go b/backend/internal/handler/admin/ops_ingress_reject_handler_test.go new file mode 100644 index 000000000..bb8949294 --- /dev/null +++ b/backend/internal/handler/admin/ops_ingress_reject_handler_test.go @@ -0,0 +1,41 @@ +package admin + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func newIngressRejectHandlerForTest() *OpsHandler { + return NewOpsHandler(service.NewOpsService(nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)) +} + +func TestListIngressRejectsValidatesFilters(t *testing.T) { + gin.SetMode(gin.TestMode) + for _, path := range []string{ + "/api/v1/admin/ops/ingress-rejections?reason=not-valid", + "/api/v1/admin/ops/ingress-rejections?client_ip=not-an-ip", + "/api/v1/admin/ops/ingress-rejections?user_id=0", + "/api/v1/admin/ops/ingress-rejections?api_key_id=-1", + } { + t.Run(path, func(t *testing.T) { + context, _ := gin.CreateTestContext(httptest.NewRecorder()) + context.Request = httptest.NewRequest(http.MethodGet, path, nil) + newIngressRejectHandlerForTest().ListIngressRejects(context) + require.Equal(t, http.StatusBadRequest, context.Writer.Status()) + }) + } +} + +func TestGetIngressRejectHealthAvailable(t *testing.T) { + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Request = httptest.NewRequest(http.MethodGet, "/api/v1/admin/ops/ingress-rejections/health", nil) + newIngressRejectHandlerForTest().GetIngressRejectHealth(context) + require.Equal(t, http.StatusOK, recorder.Code) + require.Contains(t, recorder.Body.String(), `"capacity":8192`) +} diff --git a/backend/internal/handler/admin/setting_handler_update.go b/backend/internal/handler/admin/setting_handler_update.go index 1e11898e0..40343ed31 100644 --- a/backend/internal/handler/admin/setting_handler_update.go +++ b/backend/internal/handler/admin/setting_handler_update.go @@ -1688,6 +1688,9 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) { response.ErrorFrom(c, err) return } + if h.opsService != nil { + h.opsService.SetMonitoringEnabled(settings.OpsMonitoringEnabled) + } // Update OpenAI fast policy (stored under dedicated key, only when provided). if req.OpenAIFastPolicySettings != nil { diff --git a/backend/internal/handler/gateway_handler.go b/backend/internal/handler/gateway_handler.go index 3f7c9ff72..f5d99f3c3 100644 --- a/backend/internal/handler/gateway_handler.go +++ b/backend/internal/handler/gateway_handler.go @@ -166,6 +166,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) { h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body") return } + body = parsedReq.Body.Bytes() reqModel := parsedReq.Model reqStream := parsedReq.Stream reqLog = reqLog.With(zap.String("model", reqModel), zap.Bool("stream", reqStream)) @@ -1882,6 +1883,7 @@ func (h *GatewayHandler) CountTokens(c *gin.Context) { h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body") return } + body = parsedReq.Body.Bytes() // count_tokens 走 messages 严格校验时,复用已解析请求,避免二次反序列化。 SetClaudeCodeClientContext(c, body, parsedReq) reqLog = reqLog.With(zap.String("model", parsedReq.Model), zap.Bool("stream", parsedReq.Stream)) diff --git a/backend/internal/handler/grok_media.go b/backend/internal/handler/grok_media.go index 98f1ee69b..d6c791fca 100644 --- a/backend/internal/handler/grok_media.go +++ b/backend/internal/handler/grok_media.go @@ -168,6 +168,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. sameAccountRetryCount := make(map[int64]int) var lastFailoverErr *service.UpstreamFailoverError var oauth429FailoverState service.OpenAIOAuth429FailoverState + mediaEligibilityRejected := false switchCount := 0 maxAccountSwitches := h.maxAccountSwitches if maxAccountSwitches <= 0 { @@ -203,7 +204,8 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. zap.Error(err), zap.Int("excluded_account_count", len(failedAccountIDs)), ) - if endpoint.IsGenerationRequest() && len(failedAccountIDs) == 0 && errors.Is(err, service.ErrNoAvailableAccounts) { + if endpoint.IsGenerationRequest() && errors.Is(err, service.ErrNoAvailableAccounts) && + (len(failedAccountIDs) == 0 || (mediaEligibilityRejected && lastFailoverErr == nil)) { markOpsRoutingCapacityLimitedIfNoAvailable(c, err) h.errorResponse(c, http.StatusServiceUnavailable, "grok_media_no_eligible_account", "No eligible Grok media accounts") return @@ -247,6 +249,25 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. ) account := selection.Account + if endpoint.IsGenerationRequest() { + eligible, eligibilityReason, eligibilityErr := h.ensureGrokMediaAccountEligibility(requestCtx, account) + if !eligible { + 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 + } + } sessionHash = ensureOpenAIPoolModeSessionHash(sessionHash, account) setOpsSelectedAccount(c, account.ID, account.Platform) @@ -366,6 +387,20 @@ 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 diff --git a/backend/internal/handler/grok_media_test.go b/backend/internal/handler/grok_media_test.go index 8c52bb3c8..7c596c26f 100644 --- a/backend/internal/handler/grok_media_test.go +++ b/backend/internal/handler/grok_media_test.go @@ -1,12 +1,26 @@ package handler import ( + "context" + "errors" "testing" "github.com/Wei-Shaw/sub2api/internal/service" "github.com/stretchr/testify/require" ) +type grokMediaEligibilityProberStub struct { + eligible bool + reason string + err error + calls int +} + +func (s *grokMediaEligibilityProberStub) ProbeMediaEligibility(context.Context, int64) (bool, string, error) { + s.calls++ + return s.eligible, s.reason, s.err +} + func TestShouldRecordGrokMediaUsage(t *testing.T) { tests := []struct { name string @@ -92,3 +106,55 @@ func TestGrokMediaScheduleModelUsesNormalizedMappedUpstream(t *testing.T) { require.Equal(t, "mapped-video-model", grokMediaScheduleModel(account, "grok-imagine-video", &service.OpenAIForwardResult{})) require.Equal(t, "grok-imagine-video", grokMediaScheduleModel(nil, " grok-imagine-video ", nil)) } + +func TestEnsureGrokMediaAccountEligibility(t *testing.T) { + t.Run("non oauth account does not probe", func(t *testing.T) { + prober := &grokMediaEligibilityProberStub{} + h := &OpenAIGatewayHandler{grokMediaEligibilityProber: prober} + account := &service.Account{Platform: service.PlatformGrok, Type: service.AccountTypeAPIKey} + + eligible, reason, err := h.ensureGrokMediaAccountEligibility(context.Background(), account) + + require.NoError(t, err) + require.True(t, eligible) + require.Equal(t, "non_oauth", reason) + require.Zero(t, prober.calls) + }) + + t.Run("unobserved oauth is probed before forwarding", func(t *testing.T) { + prober := &grokMediaEligibilityProberStub{eligible: true, reason: "eligible"} + h := &OpenAIGatewayHandler{grokMediaEligibilityProber: prober} + account := &service.Account{ID: 7, Platform: service.PlatformGrok, Type: service.AccountTypeOAuth} + + eligible, reason, err := h.ensureGrokMediaAccountEligibility(context.Background(), account) + + require.NoError(t, err) + require.True(t, eligible) + require.Equal(t, "eligible", reason) + require.Equal(t, 1, prober.calls) + }) + + t.Run("missing prober fails closed", func(t *testing.T) { + h := &OpenAIGatewayHandler{} + account := &service.Account{ID: 8, Platform: service.PlatformGrok, Type: service.AccountTypeOAuth} + + eligible, reason, err := h.ensureGrokMediaAccountEligibility(context.Background(), account) + + require.Error(t, err) + require.False(t, eligible) + require.Equal(t, "billing_probe_unavailable", reason) + }) + + t.Run("probe failure fails closed", func(t *testing.T) { + probeErr := errors.New("probe failed") + prober := &grokMediaEligibilityProberStub{reason: "billing_unobserved", err: probeErr} + h := &OpenAIGatewayHandler{grokMediaEligibilityProber: prober} + account := &service.Account{ID: 9, Platform: service.PlatformGrok, Type: service.AccountTypeOAuth} + + eligible, reason, err := h.ensureGrokMediaAccountEligibility(context.Background(), account) + + require.ErrorIs(t, err, probeErr) + require.False(t, eligible) + require.Equal(t, "billing_unobserved", reason) + }) +} diff --git a/backend/internal/handler/no_account_error.go b/backend/internal/handler/no_account_error.go index 001cef611..842935aa2 100644 --- a/backend/internal/handler/no_account_error.go +++ b/backend/internal/handler/no_account_error.go @@ -38,11 +38,12 @@ type noAccountErrorClassification struct { // The classifier intentionally does not consume the original error: the // selection layer never tells us *why* the pool came up empty (rate-limited // vs. unsupported model are both wrapped as ErrNoAvailableAccounts). Instead -// we re-check pool composition through DiagnoseModelAvailabilityForPlatform, -// which only inspects model_mapping configuration and ignores transient -// state. That guarantees a 404 is only returned when no operator action -// short of editing the account's model_mapping could make this request -// succeed. +// we re-check pool composition through DiagnoseModelAvailabilityForPlatform. +// Its dedicated database query considers only persistent eligibility +// (active status + schedulable setting) and model_mapping, bypassing scheduler +// snapshots and transient filters. That guarantees a 404 is only returned +// when persistent account/group/model configuration must change before the +// request can succeed. // // routingModel is the model name that account selection actually compared // against (i.e. after group-level dispatch mapping). displayModel is the diff --git a/backend/internal/handler/no_account_error_test.go b/backend/internal/handler/no_account_error_test.go index 174da82cc..8c6784e84 100644 --- a/backend/internal/handler/no_account_error_test.go +++ b/backend/internal/handler/no_account_error_test.go @@ -154,6 +154,20 @@ func TestClassifyNoAccountError_HasModelSupport_KeepsRoutingMessageGenerationToC require.False(t, cls.ModelNotFound) } +func TestClassifyNoAccountError_ModelSupportedOnlyByRateLimitedAccount_Returns503(t *testing.T) { + c := newTestGinContextWithRequest() + // The diagnoser's configured-state lookup still sees the model-supporting + // account even though normal scheduling has excluded it during cooldown. + fd := &fakeDiagnoser{resp: service.ModelAvailabilityDiagnosis{HasAccountsInPool: true, HasModelSupport: true}} + apiKey := &service.APIKey{GroupID: ptrInt64(7)} + + cls := classifyNoAccountErrorFromGin(c, fd, apiKey, "claude-opus-4-8", "claude-opus-4-8", service.PlatformAnthropic) + + require.Equal(t, http.StatusServiceUnavailable, cls.Status) + require.Equal(t, "api_error", cls.ErrType) + require.False(t, cls.ModelNotFound, "temporary account cooldown must remain retryable") +} + func TestClassifyNoAccountError_NoAccountsInPool_Stays503(t *testing.T) { c := newTestGinContextWithRequest() fd := &fakeDiagnoser{resp: service.ModelAvailabilityDiagnosis{HasAccountsInPool: false, HasModelSupport: false}} diff --git a/backend/internal/handler/openai_chat_completions.go b/backend/internal/handler/openai_chat_completions.go index 974d35ded..e16962ba5 100644 --- a/backend/internal/handler/openai_chat_completions.go +++ b/backend/internal/handler/openai_chat_completions.go @@ -305,7 +305,10 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) { upstreamErrorAlreadyCommunicated := openAIForwardErrorAlreadyCommunicated(c, writerSizeBeforeForward, err) wroteFallback := false if !upstreamErrorAlreadyCommunicated { - wroteFallback = h.ensureForwardErrorResponse(c, streamStarted) + wroteFallback = h.ensureOpenAIStreamReadErrorResponse(c, err, streamStarted) + if !wroteFallback { + wroteFallback = h.ensureForwardErrorResponse(c, streamStarted) + } } reqLog.Warn("openai_chat_completions.forward_failed", zap.Int64("account_id", account.ID), diff --git a/backend/internal/handler/openai_gateway_count_tokens.go b/backend/internal/handler/openai_gateway_count_tokens.go index 114337460..a9c2710b9 100644 --- a/backend/internal/handler/openai_gateway_count_tokens.go +++ b/backend/internal/handler/openai_gateway_count_tokens.go @@ -67,6 +67,7 @@ func (h *OpenAIGatewayHandler) CountTokens(c *gin.Context) { h.anthropicErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body") return } + body = parsedReq.Body.Bytes() if parsedReq.Model == "" { h.anthropicErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "model is required") return diff --git a/backend/internal/handler/openai_gateway_credential_failover_test.go b/backend/internal/handler/openai_gateway_credential_failover_test.go index 07a5e4999..d650ef6a2 100644 --- a/backend/internal/handler/openai_gateway_credential_failover_test.go +++ b/backend/internal/handler/openai_gateway_credential_failover_test.go @@ -206,6 +206,10 @@ func TestOpsRecoveredCredentialFailoverUsesAccountAuthAttribution(t *testing.T) require.NotContains(t, job.entry.ErrorMessage, "earlier inference failure") require.NotNil(t, job.entry.UpstreamStatusCode) require.Zero(t, *job.entry.UpstreamStatusCode) - require.Len(t, job.entry.UpstreamErrors, 2) - require.Equal(t, http.StatusForbidden, job.entry.UpstreamErrors[0].UpstreamStatusCode) + require.Nil(t, job.entry.UpstreamErrors) + require.NotNil(t, job.entry.UpstreamErrorsJSON) + events, err := service.ParseOpsUpstreamErrors(*job.entry.UpstreamErrorsJSON) + require.NoError(t, err) + require.Len(t, events, 2) + require.Equal(t, http.StatusForbidden, events[0].UpstreamStatusCode) } diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index ec4e06dd8..602b971f4 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -28,18 +28,23 @@ 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 - securityAuditCoordinator *securityaudit.Coordinator - opsService *service.OpsService - concurrencyHelper *ConcurrencyHelper - imageLimiter *imageConcurrencyLimiter - maxAccountSwitches int - cfg *config.Config + gatewayService *service.OpenAIGatewayService + billingCacheService *service.BillingCacheService + apiKeyService *service.APIKeyService + usageRecordWorkerPool *service.UsageRecordWorkerPool + errorPassthroughService *service.ErrorPassthroughService + contentModerationService *service.ContentModerationService + securityAuditCoordinator *securityaudit.Coordinator + grokMediaEligibilityProber grokMediaEligibilityProber + opsService *service.OpsService + concurrencyHelper *ConcurrencyHelper + imageLimiter *imageConcurrencyLimiter + maxAccountSwitches int + cfg *config.Config +} + +type grokMediaEligibilityProber interface { + ProbeMediaEligibility(ctx context.Context, accountID int64) (bool, string, error) } const maxOpenAIFirstOutputTimeoutSwitches = 1 @@ -2231,12 +2236,29 @@ func (h *OpenAIGatewayHandler) mapUpstreamError(statusCode int) (int, string, st // handleStreamingAwareError handles errors that may occur after streaming has started func (h *OpenAIGatewayHandler) handleStreamingAwareError(c *gin.Context, status int, errType, message string, streamStarted bool) { + h.handleStreamingAwareErrorWithCode(c, status, errType, "", message, streamStarted, false) +} + +func (h *OpenAIGatewayHandler) handleStreamingAwareErrorWithCode( + c *gin.Context, + status int, + errType string, + code string, + message string, + streamStarted bool, + countTowardsSLA bool, +) { // body-signal compact 心跳可能已把响应头提交为 200:先停心跳(建立 // happens-before,接管 ResponseWriter),并升级为流内错误处理。 if service.StopOpenAICompactSSEKeepaliveCommitted(c) { streamStarted = true } if streamStarted { + if countTowardsSLA { + service.MarkOpsStreamFailure(c, errType, code, message, status) + } else { + service.MarkOpsStreamError(c, errType, message, status) + } // /v1/responses 的严格 SDK(Codex CLI)要求终止事件必须属于 // response.completed/failed/incomplete/cancelled 集合。 // 通用 `event: error` 帧不被识别为终止事件,会导致 @@ -2249,8 +2271,15 @@ func (h *OpenAIGatewayHandler) handleStreamingAwareError(c *gin.Context, status // Stream already started, send error as SSE event then close flusher, ok := c.Writer.(http.Flusher) if ok { - // SSE 错误事件固定 schema,使用 Quote 直拼可避免额外 Marshal 分配。 - errorEvent := "event: error\ndata: " + `{"error":{"type":` + strconv.Quote(errType) + `,"message":` + strconv.Quote(message) + `}}` + "\n\n" + errorObject := gin.H{"type": errType, "message": message} + if code != "" { + errorObject["code"] = code + } + payload, err := json.Marshal(gin.H{"error": errorObject}) + if err != nil { + payload = []byte(`{"error":{"type":"upstream_error","message":"Upstream request failed"}}`) + } + errorEvent := "event: error\ndata: " + string(payload) + "\n\n" if _, err := fmt.Fprint(c.Writer, errorEvent); err != nil { _ = c.Error(err) } @@ -2260,7 +2289,27 @@ func (h *OpenAIGatewayHandler) handleStreamingAwareError(c *gin.Context, status } // Normal case: return JSON response with proper status code - h.errorResponse(c, status, errType, message) + if code == "" { + h.errorResponse(c, status, errType, message) + return + } + c.JSON(status, gin.H{"error": gin.H{ + "type": errType, "code": code, "message": message, + }}) +} + +func (h *OpenAIGatewayHandler) ensureOpenAIStreamReadErrorResponse(c *gin.Context, err error, streamStarted bool) bool { + code, message, ok := service.OpenAIUpstreamStreamReadErrorDetails(err) + if !ok || c == nil || c.Writer == nil || service.IsResponseCommitted(c) { + return false + } + if c.Writer.Written() { + streamStarted = true + } + h.handleStreamingAwareErrorWithCode( + c, http.StatusBadGateway, "upstream_error", code, message, streamStarted, true, + ) + return true } // ensureForwardErrorResponse 在 Forward 返回错误但尚未写响应时补写统一错误响应。 diff --git a/backend/internal/handler/openai_gateway_handler_test.go b/backend/internal/handler/openai_gateway_handler_test.go index 931f3b3f6..635c22c3f 100644 --- a/backend/internal/handler/openai_gateway_handler_test.go +++ b/backend/internal/handler/openai_gateway_handler_test.go @@ -96,6 +96,35 @@ func TestOpenAIHandleStreamingAwareError_JSONEscaping(t *testing.T) { } } +func TestOpenAIHandleStreamingAwareErrorWithCode_EmitsStableClassification(t *testing.T) { + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + + h := &OpenAIGatewayHandler{} + h.handleStreamingAwareErrorWithCode( + c, + http.StatusBadGateway, + "upstream_error", + service.OpenAIUpstreamHTTP2StreamErrorCode, + "Upstream HTTP/2 stream failed", + true, + true, + ) + + body := w.Body.String() + require.Contains(t, body, "event: error\n") + require.Equal(t, "upstream_error", gjson.Get(body[strings.Index(body, "{"):], "error.type").String()) + require.Equal(t, service.OpenAIUpstreamHTTP2StreamErrorCode, gjson.Get(body[strings.Index(body, "{"):], "error.code").String()) + require.NotContains(t, body, "stream ID") + + streamErr, ok := service.GetOpsStreamError(c) + require.True(t, ok) + require.True(t, streamErr.CountTowardsSLA) + require.Equal(t, http.StatusBadGateway, streamErr.IntendedStatus) +} + func TestOpenAIForwardSucceededForScheduling(t *testing.T) { require.True(t, openAIForwardSucceededForScheduling(nil)) require.True(t, openAIForwardSucceededForScheduling(&service.OpenAIForwardResult{})) @@ -1875,6 +1904,212 @@ func TestOpenAIResponsesWebSocket_FailoverOnUpstreamUsageLimitEvent(t *testing.T require.Equal(t, []int64{int64(9902)}, accountRepo.rateLimitedIDs) } +func TestOpenAIResponsesWebSocket_FirstOutputTimeoutWithoutDownstreamReusesClientForOneFailover(t *testing.T) { + gin.SetMode(gin.TestMode) + + firstHitCh := make(chan []byte, 1) + secondHitCh := make(chan []byte, 1) + var firstConnections atomic.Int32 + var secondConnections atomic.Int32 + + firstUpstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + firstConnections.Add(1) + conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover}) + if err != nil { + return + } + defer func() { _ = conn.CloseNow() }() + + readCtx, cancelRead := context.WithTimeout(r.Context(), 3*time.Second) + _, payload, readErr := conn.Read(readCtx) + cancelRead() + if readErr == nil { + firstHitCh <- payload + } + + select { + case <-r.Context().Done(): + case <-time.After(3 * time.Second): + } + })) + defer firstUpstream.Close() + + secondUpstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + secondConnections.Add(1) + conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover}) + if err != nil { + return + } + defer func() { _ = conn.CloseNow() }() + + readCtx, cancelRead := context.WithTimeout(r.Context(), 3*time.Second) + _, payload, readErr := conn.Read(readCtx) + cancelRead() + if readErr == nil { + secondHitCh <- payload + } + + for _, event := range []string{ + `{"type":"response.created","response":{"id":"resp_ws_timeout_b","model":"gpt-5.1"}}`, + `{"type":"response.output_text.delta","response_id":"resp_ws_timeout_b","delta":"recovered"}`, + `{"type":"response.completed","response":{"id":"resp_ws_timeout_b","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`, + } { + writeCtx, cancelWrite := context.WithTimeout(r.Context(), 3*time.Second) + writeErr := conn.Write(writeCtx, coderws.MessageText, []byte(event)) + cancelWrite() + if writeErr != nil { + return + } + } + readCtx, cancelRead = context.WithTimeout(r.Context(), 3*time.Second) + _, _, _ = conn.Read(readCtx) + cancelRead() + })) + defer secondUpstream.Close() + + groupID := int64(4212) + accounts := []service.Account{ + { + ID: 9912, + Name: "openai-ws-first-semantic-timeout", + Platform: service.PlatformOpenAI, + Type: service.AccountTypeAPIKey, + Status: service.StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 1, + Credentials: map[string]any{"api_key": "sk-first", "base_url": firstUpstream.URL}, + Extra: map[string]any{ + "openai_apikey_responses_websockets_v2_enabled": true, + "openai_apikey_responses_websockets_v2_mode": service.OpenAIWSIngressModePassthrough, + }, + }, + { + ID: 9913, + Name: "openai-ws-failover-healthy", + Platform: service.PlatformOpenAI, + Type: service.AccountTypeAPIKey, + Status: service.StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 2, + Credentials: map[string]any{"api_key": "sk-second", "base_url": secondUpstream.URL}, + Extra: map[string]any{ + "openai_apikey_responses_websockets_v2_enabled": true, + "openai_apikey_responses_websockets_v2_mode": service.OpenAIWSIngressModePassthrough, + }, + }, + } + + cfg := &config.Config{} + cfg.RunMode = config.RunModeSimple + cfg.Default.RateMultiplier = 1 + cfg.Security.URLAllowlist.Enabled = false + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + cfg.Gateway.OpenAIFirstOutputTimeoutSeconds = 1 + cfg.Gateway.OpenAIWS.Enabled = true + cfg.Gateway.OpenAIWS.APIKeyEnabled = true + cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true + cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true + cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.IngressInterTurnIdleTimeoutSeconds = 3 + cfg.Gateway.MaxAccountSwitches = 3 + + accountRepo := &openAIWSFailoverHandlerAccountRepoStub{accounts: accounts} + rateLimitSvc := service.NewRateLimitService(accountRepo, nil, cfg, nil, nil) + billingCacheSvc := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil) + gatewaySvc := service.NewOpenAIGatewayService( + accountRepo, nil, nil, nil, nil, nil, nil, cfg, nil, nil, + service.NewBillingService(cfg, nil), rateLimitSvc, billingCacheSvc, + nil, &service.DeferredService{}, nil, nil, nil, nil, nil, nil, nil, + ) + cache := &concurrencyCacheMock{ + acquireUserSlotFn: func(context.Context, int64, int, string) (bool, error) { return true, nil }, + acquireAccountSlotFn: func(context.Context, int64, int, string) (bool, error) { + return true, nil + }, + } + h := &OpenAIGatewayHandler{ + gatewayService: gatewaySvc, + billingCacheService: billingCacheSvc, + apiKeyService: &service.APIKeyService{}, + concurrencyHelper: NewConcurrencyHelper(service.NewConcurrencyService(cache), SSEPingFormatNone, time.Second), + maxAccountSwitches: 3, + } + + apiKey := &service.APIKey{ + ID: 1812, + GroupID: &groupID, + User: &service.User{ID: 1712, Status: service.StatusActive}, + Group: &service.Group{ID: groupID, Platform: service.PlatformOpenAI, Status: service.StatusActive}, + } + handlerDone := make(chan struct{}) + router := gin.New() + router.Use(func(c *gin.Context) { + c.Set(string(middleware.ContextKeyAPIKey), apiKey) + c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: apiKey.User.ID, Concurrency: 1}) + c.Next() + }) + router.GET("/openai/v1/responses", func(c *gin.Context) { + h.ResponsesWebSocket(c) + close(handlerDone) + }) + handlerServer := httptest.NewServer(router) + defer handlerServer.Close() + + dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second) + clientConn, _, err := coderws.Dial( + dialCtx, + "ws"+strings.TrimPrefix(handlerServer.URL, "http")+"/openai/v1/responses", + &coderws.DialOptions{CompressionMode: coderws.CompressionContextTakeover}, + ) + cancelDial() + require.NoError(t, err) + defer func() { _ = clientConn.CloseNow() }() + + writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second) + err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1","stream":false}`)) + cancelWrite() + require.NoError(t, err) + + var eventTypes []string + readCtx, cancelRead := context.WithTimeout(context.Background(), 6*time.Second) + for { + _, event, readErr := clientConn.Read(readCtx) + require.NoError(t, readErr) + eventType := gjson.GetBytes(event, "type").String() + eventTypes = append(eventTypes, eventType) + if eventType == "response.completed" { + require.Equal(t, "resp_ws_timeout_b", gjson.GetBytes(event, "response.id").String()) + break + } + } + cancelRead() + require.Contains(t, eventTypes, "response.output_text.delta") + require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done")) + + select { + case <-handlerDone: + case <-time.After(3 * time.Second): + t.Fatal("websocket handler did not finish after healthy failover turn") + } + select { + case <-firstHitCh: + case <-time.After(3 * time.Second): + t.Fatal("first upstream did not receive replayable request") + } + select { + case <-secondHitCh: + case <-time.After(3 * time.Second): + t.Fatal("second upstream did not receive replayed request") + } + require.Equal(t, int32(1), firstConnections.Load()) + require.Equal(t, int32(1), secondConnections.Load()) + require.NotContains(t, accountRepo.rateLimitedIDs, int64(9913), "healthy failover account must not be penalized") +} + func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSUsageLogCase) openAIResponsesWSUsageLogResult { t.Helper() gin.SetMode(gin.TestMode) diff --git a/backend/internal/handler/ops_error_logger.go b/backend/internal/handler/ops_error_logger.go index f32ca9c05..6f09e6408 100644 --- a/backend/internal/handler/ops_error_logger.go +++ b/backend/internal/handler/ops_error_logger.go @@ -72,24 +72,10 @@ const ( opsErrorLogMinQueueSize = 256 opsErrorLogMaxQueueSize = 8192 opsErrorLogBatchSize = 32 + opsErrorLogMaxQueueBytes = 32 * 1024 * 1024 + opsErrorLogMaxUserAgentBytes = 512 ) -// looksLikeSystemKey 粗筛"形似本系统 key"的输入:长度 16-128 且仅含 [a-zA-Z0-9_-]。 -// 不用前缀匹配(APIKeyPrefix 可配置)。用于反查审计表前挡掉随机扫描的乱码输入。 -func looksLikeSystemKey(key string) bool { - if len(key) < 16 || len(key) > 128 { - return false - } - for _, c := range key { - allowed := (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || - (c >= '0' && c <= '9') || c == '_' || c == '-' - if !allowed { - return false - } - } - return true -} - // keyPrefix 返回脱敏前缀(前 n 个字符);不足 n 则原样返回。 func keyPrefix(key string, n int) string { if len(key) <= n { @@ -98,43 +84,26 @@ func keyPrefix(key string, n int) string { return key[:n] } -// extractAttemptedKey 按认证中间件同样的顺序从请求头提取提交的 key 明文。 -// 与 api_key_auth.go:43-59 一致:Authorization 仅取 Bearer scheme,非 Bearer 则忽略并继续 x-api-key → x-goog-api-key。 -func extractAttemptedKey(c *gin.Context) string { - if h := c.GetHeader("Authorization"); h != "" { - parts := strings.SplitN(h, " ", 2) - if len(parts) == 2 && strings.EqualFold(parts[0], "Bearer") { - return strings.TrimSpace(parts[1]) - } - // 非 Bearer:与中间件一致,忽略 Authorization,继续尝试其它 header(不在此 return)。 - } - if k := c.GetHeader("x-api-key"); k != "" { - return strings.TrimSpace(k) - } - if k := c.GetHeader("x-goog-api-key"); k != "" { - return strings.TrimSpace(k) - } - return "" -} - type opsErrorLogJob struct { - ops *service.OpsService - entry *service.OpsInsertErrorLogInput + ops *service.OpsService + entry *service.OpsInsertErrorLogInput + queuedBytes int64 } var ( opsErrorLogOnce sync.Once opsErrorLogQueue chan opsErrorLogJob - opsErrorLogStopOnce sync.Once - opsErrorLogWorkersWg sync.WaitGroup - opsErrorLogMu sync.RWMutex - opsErrorLogStopping bool - opsErrorLogQueueLen atomic.Int64 - opsErrorLogEnqueued atomic.Int64 - opsErrorLogDropped atomic.Int64 - opsErrorLogProcessed atomic.Int64 - opsErrorLogSanitized atomic.Int64 + opsErrorLogStopOnce sync.Once + opsErrorLogWorkersWg sync.WaitGroup + opsErrorLogMu sync.RWMutex + opsErrorLogStopping bool + opsErrorLogQueueLen atomic.Int64 + opsErrorLogQueueBytes atomic.Int64 + opsErrorLogEnqueued atomic.Int64 + opsErrorLogDropped atomic.Int64 + opsErrorLogProcessed atomic.Int64 + opsErrorLogSanitized atomic.Int64 opsErrorLogLastDropLogAt atomic.Int64 @@ -154,6 +123,7 @@ func startOpsErrorLogWorkers() { workerCount, queueSize := opsErrorLogConfig() opsErrorLogQueue = make(chan opsErrorLogJob, queueSize) opsErrorLogQueueLen.Store(0) + opsErrorLogQueueBytes.Store(0) opsErrorLogWorkersWg.Add(workerCount) for i := 0; i < workerCount; i++ { @@ -165,6 +135,7 @@ func startOpsErrorLogWorkers() { return } opsErrorLogQueueLen.Add(-1) + opsErrorLogQueueBytes.Add(-job.queuedBytes) batch := make([]opsErrorLogJob, 0, opsErrorLogBatchSize) batch = append(batch, job) @@ -184,6 +155,7 @@ func startOpsErrorLogWorkers() { return } opsErrorLogQueueLen.Add(-1) + opsErrorLogQueueBytes.Add(-nextJob.queuedBytes) batch = append(batch, nextJob) case <-timer.C: break batchLoop @@ -239,6 +211,20 @@ func enqueueOpsErrorLog(ops *service.OpsService, entry *service.OpsInsertErrorLo if ops == nil || entry == nil { return } + entry.UserAgent = normalizeOpsPersistentUserAgent(entry.UserAgent) + if entry.ErrorBody != "" { + originalBody := entry.ErrorBody + body, truncated := service.SanitizeOpsErrorBodyForQueue(originalBody) + entry.ErrorBody = body + if truncated || body != originalBody { + opsErrorLogSanitized.Add(1) + } + } + if err := service.SanitizeOpsUpstreamErrorsForQueue(entry); err != nil { + opsErrorLogDropped.Add(1) + maybeLogOpsErrorLogDrop() + return + } select { case <-opsErrorLogShutdownCh: return @@ -259,18 +245,29 @@ func enqueueOpsErrorLog(ops *service.OpsService, entry *service.OpsInsertErrorLo if opsErrorLogStopping || opsErrorLogQueue == nil { return } + queuedBytes := estimateOpsErrorLogJobBytes(entry) + if !reserveOpsErrorLogQueueBytes(queuedBytes) { + opsErrorLogDropped.Add(1) + maybeLogOpsErrorLogDrop() + return + } select { - case opsErrorLogQueue <- opsErrorLogJob{ops: ops, entry: entry}: - opsErrorLogQueueLen.Add(1) + case opsErrorLogQueue <- opsErrorLogJob{ops: ops, entry: entry, queuedBytes: queuedBytes}: opsErrorLogEnqueued.Add(1) default: + opsErrorLogQueueLen.Add(-1) + opsErrorLogQueueBytes.Add(-queuedBytes) // Queue is full; drop to avoid blocking request handling. opsErrorLogDropped.Add(1) maybeLogOpsErrorLogDrop() } } +func normalizeOpsPersistentUserAgent(value string) string { + return truncateString(strings.TrimSpace(strings.ToValidUTF8(value, "")), opsErrorLogMaxUserAgentBytes) +} + func StopOpsErrorLogWorkers() bool { opsErrorLogStopOnce.Do(func() { opsErrorLogShutdownOnce.Do(func() { @@ -293,6 +290,7 @@ func stopOpsErrorLogWorkers() bool { if ch == nil { opsErrorLogQueueLen.Store(0) + opsErrorLogQueueBytes.Store(0) return true } @@ -305,6 +303,7 @@ func stopOpsErrorLogWorkers() bool { select { case <-done: opsErrorLogQueueLen.Store(0) + opsErrorLogQueueBytes.Store(0) return true case <-time.After(opsErrorLogDrainTimeout): return false @@ -315,6 +314,14 @@ func OpsErrorLogQueueLength() int64 { return opsErrorLogQueueLen.Load() } +func OpsErrorLogQueueBytes() int64 { + return opsErrorLogQueueBytes.Load() +} + +func OpsErrorLogQueueBytesCapacity() int64 { + return opsErrorLogMaxQueueBytes +} + func OpsErrorLogQueueCapacity() int { opsErrorLogMu.RLock() ch := opsErrorLogQueue @@ -355,12 +362,15 @@ func maybeLogOpsErrorLogDrop() { } queued := opsErrorLogQueueLen.Load() + queuedBytes := opsErrorLogQueueBytes.Load() queueCap := OpsErrorLogQueueCapacity() log.Printf( - "[OpsErrorLogger] queue is full; dropping logs (queued=%d cap=%d enqueued_total=%d dropped_total=%d processed_total=%d sanitized_total=%d)", + "[OpsErrorLogger] queue is full; dropping logs (queued=%d cap=%d queued_bytes=%d bytes_cap=%d enqueued_total=%d dropped_total=%d processed_total=%d sanitized_total=%d)", queued, queueCap, + queuedBytes, + opsErrorLogMaxQueueBytes, opsErrorLogEnqueued.Load(), opsErrorLogDropped.Load(), opsErrorLogProcessed.Load(), @@ -368,6 +378,46 @@ func maybeLogOpsErrorLogDrop() { ) } +func reserveOpsErrorLogQueueBytes(size int64) bool { + if size < 1 { + size = 1 + } + for { + current := opsErrorLogQueueBytes.Load() + if current > opsErrorLogMaxQueueBytes-size { + return false + } + if opsErrorLogQueueBytes.CompareAndSwap(current, current+size) { + opsErrorLogQueueLen.Add(1) + return true + } + } +} + +func estimateOpsErrorLogJobBytes(entry *service.OpsInsertErrorLogInput) int64 { + if entry == nil { + return 1 + } + const fixedOverhead = 512 + size := fixedOverhead + len(entry.RequestID) + len(entry.ClientRequestID) + + len(entry.Platform) + len(entry.Model) + len(entry.RequestPath) + + len(entry.InboundEndpoint) + len(entry.UpstreamEndpoint) + + len(entry.RequestedModel) + len(entry.UpstreamModel) + len(entry.UserAgent) + + len(entry.ErrorPhase) + len(entry.ErrorType) + len(entry.Severity) + + len(entry.ErrorMessage) + len(entry.ErrorBody) + len(entry.ErrorSource) + + len(entry.ErrorOwner) + len(entry.APIKeyPrefix) + if entry.UpstreamErrorMessage != nil { + size += len(*entry.UpstreamErrorMessage) + } + if entry.UpstreamErrorDetail != nil { + size += len(*entry.UpstreamErrorDetail) + } + if entry.UpstreamErrorsJSON != nil { + size += len(*entry.UpstreamErrorsJSON) + } + return int64(size) +} + func opsErrorLogConfig() (workerCount int, queueSize int) { workerCount = runtime.GOMAXPROCS(0) * 2 if workerCount < opsErrorLogMinWorkerCount { @@ -470,9 +520,12 @@ type opsCaptureWriter struct { gin.ResponseWriter limit int buf bytes.Buffer + ctx *gin.Context } -const opsCaptureWriterLimit = 64 * 1024 +const opsCaptureWriterLimit = service.OpsErrorLogQueueBodyMaxBytes + +const opsCaptureWriterPoolMaxRetainedCapacity = service.OpsErrorLogQueueBodyMaxBytes var opsCaptureWriterPool = sync.Pool{ New: func() any { @@ -496,11 +549,19 @@ func releaseOpsCaptureWriter(w *opsCaptureWriter) { return } w.ResponseWriter = nil + w.ctx = nil w.limit = opsCaptureWriterLimit + if !shouldPoolOpsCaptureWriter(w) { + return + } w.buf.Reset() opsCaptureWriterPool.Put(w) } +func shouldPoolOpsCaptureWriter(w *opsCaptureWriter) bool { + return w != nil && w.buf.Cap() <= opsCaptureWriterPoolMaxRetainedCapacity +} + func (w *opsCaptureWriter) Status() int { if w.ResponseWriter == nil { return 0 @@ -577,7 +638,7 @@ func (w *opsCaptureWriter) Write(b []byte) (int, error) { if w.ResponseWriter == nil { return 0, nil } - if w.Status() >= 400 && w.limit > 0 && w.buf.Len() < w.limit { + if w.shouldCapture() && w.Status() >= 400 && w.limit > 0 && w.buf.Len() < w.limit { remaining := w.limit - w.buf.Len() if len(b) > remaining { _, _ = w.buf.Write(b[:remaining]) @@ -592,7 +653,7 @@ func (w *opsCaptureWriter) WriteString(s string) (int, error) { if w.ResponseWriter == nil { return 0, nil } - if w.Status() >= 400 && w.limit > 0 && w.buf.Len() < w.limit { + if w.shouldCapture() && w.Status() >= 400 && w.limit > 0 && w.buf.Len() < w.limit { remaining := w.limit - w.buf.Len() if len(s) > remaining { _, _ = w.buf.WriteString(s[:remaining]) @@ -603,6 +664,14 @@ func (w *opsCaptureWriter) WriteString(s string) (int, error) { return w.ResponseWriter.WriteString(s) } +func (w *opsCaptureWriter) shouldCapture() bool { + if w.ctx == nil { + return true + } + _, rejected := middleware2.GetIngressRejectReason(w.ctx) + return !rejected +} + // OpsErrorLoggerMiddleware records error responses (status >= 400) into ops_error_logs. // // Notes: @@ -612,6 +681,7 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc { return func(c *gin.Context) { originalWriter := c.Writer w := acquireOpsCaptureWriter(originalWriter) + w.ctx = c defer func() { // Restore the original writer before returning so outer middlewares // don't observe a pooled wrapper that has been released. @@ -623,6 +693,10 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc { c.Writer = w c.Next() + if _, rejected := middleware2.GetIngressRejectReason(c); rejected { + return + } + if ops == nil { return } @@ -989,7 +1063,7 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc { IsCountTokens: isCountTokensRequest(c), ErrorMessage: parsed.Message, - // Keep the full captured error body (capture is already capped at 64KB) so the + // Keep the captured error body (already capped at the queue-safe limit) so the // service layer can sanitize JSON before truncating for storage. ErrorBody: string(body), ErrorSource: errorSource, @@ -1002,7 +1076,7 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc { if apiKey != nil { entry.APIKeyID = &apiKey.ID - // 有效(未删除)key 报错时快照前缀,key 之后被删也保留;与 INVALID_API_KEY 的 attempted_key_prefix 互斥。 + // 有效 key 报错时快照前缀,key 之后被删也保留。 entry.APIKeyPrefix = keyPrefix(apiKey.Key, 8) if apiKey.User != nil { entry.UserID = &apiKey.User.ID @@ -1022,22 +1096,6 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc { entry.ClientIP = &clientIP } - // 已删除 key 归因:仅 INVALID_API_KEY 才尝试。响应已写出,此处不阻塞客户端。 - if parsed.Code == opsCodeInvalidAPIKey { - if attemptedKey := extractAttemptedKey(c); attemptedKey != "" { - entry.AttemptedKeyPrefix = keyPrefix(attemptedKey, 8) - if looksLikeSystemKey(attemptedKey) { - if res, lookupErr := ops.LookupDeletedKeyAudit(c.Request.Context(), attemptedKey); lookupErr != nil { - log.Printf("[OpsErrorLogger] LookupDeletedKeyAudit failed: %v", lookupErr) - } else if res != nil { - owner := res.UserID - entry.DeletedKeyOwnerUserID = &owner - entry.DeletedKeyName = res.KeyName - } - } - } - } - enqueueOpsErrorLog(ops, entry) } } @@ -1072,8 +1130,20 @@ func logOpsStreamError(c *gin.Context, ops *service.OpsService, wireStatus int) if classifyStatus <= 0 { classifyStatus = wireStatus } - normalizedType := normalizeOpsErrorType(streamErr.ErrType, "") - phase, isBusinessLimited, errorOwner, errorSource := classifyOpsErrorLog(c, normalizedType, streamErr.Message, "", classifyStatus) + normalizedType := normalizeOpsErrorType(streamErr.ErrType, streamErr.Code) + phase, isBusinessLimited, errorOwner, errorSource := classifyOpsErrorLog(c, normalizedType, streamErr.Message, streamErr.Code, classifyStatus) + recordedStatus := wireStatus + if streamErr.CountTowardsSLA && streamErr.IntendedStatus >= 400 { + recordedStatus = streamErr.IntendedStatus + } + errorBody := "" + if streamErr.Code != "" { + if payload, err := json.Marshal(gin.H{"error": gin.H{ + "type": normalizedType, "code": streamErr.Code, "message": streamErr.Message, + }}); err == nil { + errorBody = string(payload) + } + } apiKey := getOpsAPIKey(c) clientRequestID, _ := c.Request.Context().Value(ctxkey.ClientRequestID).(string) @@ -1140,12 +1210,12 @@ func logOpsStreamError(c *gin.Context, ops *service.OpsService, wireStatus int) ErrorPhase: phase, ErrorType: normalizedType, Severity: classifyOpsSeverity(normalizedType, classifyStatus), - StatusCode: wireStatus, + StatusCode: recordedStatus, IsBusinessLimited: isBusinessLimited, IsCountTokens: isCountTokensRequest(c), ErrorMessage: streamErr.Message, - ErrorBody: "", + ErrorBody: errorBody, ErrorSource: errorSource, ErrorOwner: errorOwner, @@ -1694,11 +1764,8 @@ func shouldSkipOpsErrorLog(ctx context.Context, ops *service.OpsService, message } // Get advanced settings to check filter configuration - settings, err := ops.GetOpsAdvancedSettings(ctx) - if err != nil || settings == nil { - // If we can't get settings, don't skip (fail open) - return false - } + _ = ctx + settings := ops.OpsAdvancedSettingsSnapshot() msgLower := strings.ToLower(message) bodyLower := strings.ToLower(body) diff --git a/backend/internal/handler/ops_error_logger_attribution_test.go b/backend/internal/handler/ops_error_logger_attribution_test.go index 9c68d8458..641fb8d62 100644 --- a/backend/internal/handler/ops_error_logger_attribution_test.go +++ b/backend/internal/handler/ops_error_logger_attribution_test.go @@ -1,39 +1,9 @@ package handler import ( - "net/http" - "net/http/httptest" "testing" - - "github.com/gin-gonic/gin" ) -func TestLooksLikeSystemKey(t *testing.T) { - cases := []struct { - in string - want bool - }{ - {"sk-abcdef0123456789", true}, - {"ABCdef_-0123456789", true}, - {"short", false}, - {"with space xxxxxxxxxx", false}, - {"汉字key1234567890", false}, - {"", false}, - } - for _, c := range cases { - if got := looksLikeSystemKey(c.in); got != c.want { - t.Errorf("looksLikeSystemKey(%q)=%v want %v", c.in, got, c.want) - } - } - long := make([]byte, 129) - for i := range long { - long[i] = 'a' - } - if looksLikeSystemKey(string(long)) { - t.Errorf("129-char key should be rejected") - } -} - func TestKeyPrefix(t *testing.T) { if got := keyPrefix("sk-3f2a9c7e", 8); got != "sk-3f2a9" { t.Errorf("keyPrefix=%q want %q", got, "sk-3f2a9") @@ -42,77 +12,3 @@ func TestKeyPrefix(t *testing.T) { t.Errorf("short key should be returned as-is, got %q", got) } } - -func TestExtractAttemptedKey(t *testing.T) { - gin.SetMode(gin.TestMode) - - cases := []struct { - name string - headers map[string]string - want string - }{ - { - name: "Bearer in Authorization", - headers: map[string]string{"Authorization": "Bearer sk-testkey0123456789"}, - want: "sk-testkey0123456789", - }, - { - name: "Bearer case-insensitive", - headers: map[string]string{"Authorization": "BEARER sk-testkey0123456789"}, - want: "sk-testkey0123456789", - }, - { - name: "x-api-key header", - headers: map[string]string{"x-api-key": "sk-xapikey0123456789"}, - want: "sk-xapikey0123456789", - }, - { - name: "x-goog-api-key header", - headers: map[string]string{"x-goog-api-key": "sk-goog0123456789"}, - want: "sk-goog0123456789", - }, - { - name: "Authorization takes priority over x-api-key", - headers: map[string]string{"Authorization": "Bearer sk-auth0123456789", "x-api-key": "sk-xapi0123456789"}, - want: "sk-auth0123456789", - }, - { - name: "x-api-key takes priority over x-goog-api-key", - headers: map[string]string{"x-api-key": "sk-xapi0123456789", "x-goog-api-key": "sk-goog0123456789"}, - want: "sk-xapi0123456789", - }, - { - name: "no key headers", - headers: map[string]string{}, - want: "", - }, - { - name: "Bearer with leading/trailing spaces trimmed", - headers: map[string]string{"Authorization": "Bearer sk-trimmed0123456789 "}, - want: "sk-trimmed0123456789", - }, - { - // 非 Bearer Authorization 应被忽略,继续 fall-through 到 x-api-key(与认证中间件一致) - name: "non-Bearer Authorization falls through to x-api-key", - headers: map[string]string{"Authorization": "junk-not-bearer", "x-api-key": "sk-realkey1234567"}, - want: "sk-realkey1234567", - }, - } - - for _, tc := range cases { - t.Run(tc.name, func(t *testing.T) { - rec := httptest.NewRecorder() - c, _ := gin.CreateTestContext(rec) - req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) - for k, v := range tc.headers { - req.Header.Set(k, v) - } - c.Request = req - - got := extractAttemptedKey(c) - if got != tc.want { - t.Errorf("extractAttemptedKey(%v) = %q, want %q", tc.headers, got, tc.want) - } - }) - } -} diff --git a/backend/internal/handler/ops_error_logger_test.go b/backend/internal/handler/ops_error_logger_test.go index 89bdd938c..3b778cdcf 100644 --- a/backend/internal/handler/ops_error_logger_test.go +++ b/backend/internal/handler/ops_error_logger_test.go @@ -1,10 +1,13 @@ package handler import ( + "context" "net/http" "net/http/httptest" + "strings" "sync" "testing" + "unicode/utf8" middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware" "github.com/Wei-Shaw/sub2api/internal/service" @@ -12,6 +15,82 @@ import ( "github.com/stretchr/testify/require" ) +type ingressRejectSettingRepo struct { + service.SettingRepository + getValueCalls int +} + +func (r *ingressRejectSettingRepo) GetValue(context.Context, string) (string, error) { + r.getValueCalls++ + return "", service.ErrSettingNotFound +} + +func (r *ingressRejectSettingRepo) GetMultiple(context.Context, []string) (map[string]string, error) { + r.getValueCalls++ + return map[string]string{}, nil +} + +func (r *ingressRejectSettingRepo) Set(context.Context, string, string) error { + return nil +} + +type ingressRejectOpsRepo struct { + service.OpsRepository + insertCalls int +} + +func (r *ingressRejectOpsRepo) InsertErrorLog(context.Context, *service.OpsInsertErrorLogInput) (int64, error) { + r.insertCalls++ + return 0, nil +} + +func (r *ingressRejectOpsRepo) BatchInsertErrorLogs(context.Context, []*service.OpsInsertErrorLogInput) (int64, error) { + r.insertCalls++ + return 0, nil +} + +func TestOpsErrorLogQueueByteBudget(t *testing.T) { + previousBytes := opsErrorLogQueueBytes.Load() + previousLen := opsErrorLogQueueLen.Load() + opsErrorLogQueueBytes.Store(0) + opsErrorLogQueueLen.Store(0) + t.Cleanup(func() { + opsErrorLogQueueBytes.Store(previousBytes) + opsErrorLogQueueLen.Store(previousLen) + }) + + if !reserveOpsErrorLogQueueBytes(opsErrorLogMaxQueueBytes - 1) { + t.Fatal("first reservation within byte budget should succeed") + } + if reserveOpsErrorLogQueueBytes(2) { + t.Fatal("reservation beyond byte budget should be rejected") + } + if got := OpsErrorLogQueueBytes(); got != opsErrorLogMaxQueueBytes-1 { + t.Fatalf("queued bytes = %d, want %d", got, opsErrorLogMaxQueueBytes-1) + } + if got := OpsErrorLogQueueLength(); got != 1 { + t.Fatalf("queue length = %d, want 1", got) + } +} + +func TestEstimateOpsErrorLogJobBytesIncludesVariablePayloads(t *testing.T) { + base := estimateOpsErrorLogJobBytes(&service.OpsInsertErrorLogInput{}) + message := "upstream message" + detail := "upstream detail" + events := `[{"error":"x"}]` + entry := &service.OpsInsertErrorLogInput{ + ErrorBody: strings.Repeat("x", 1024), + ErrorMessage: "client error", + UserAgent: "test-agent", + UpstreamErrorMessage: &message, + UpstreamErrorDetail: &detail, + UpstreamErrorsJSON: &events, + } + if got := estimateOpsErrorLogJobBytes(entry); got <= base+1024 { + t.Fatalf("estimated bytes = %d, expected variable payloads above %d", got, base+1024) + } +} + func resetOpsErrorLoggerStateForTest(t *testing.T) { t.Helper() @@ -119,6 +198,29 @@ func TestOpsCaptureWriterPool_ResetOnRelease(t *testing.T) { require.Zero(t, reused.buf.Len(), "writer should be reset before reuse") } +func TestOpsCaptureWriterPool_DropsLargeBuffers(t *testing.T) { + w := &opsCaptureWriter{} + w.buf.Grow(opsCaptureWriterPoolMaxRetainedCapacity + 1) + require.False(t, shouldPoolOpsCaptureWriter(w)) +} + +func TestEnqueueOpsErrorLog_SanitizesAndBoundsBodyBeforeQueue(t *testing.T) { + setupOpsErrorLogTestQueue(t, 1) + ops := service.NewOpsService(nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil) + secret := strings.Repeat("s", service.OpsErrorLogQueueBodyMaxBytes) + entry := &service.OpsInsertErrorLogInput{ + ErrorPhase: "request", + ErrorType: "api_error", + ErrorBody: `{"authorization":"Bearer ` + secret + `","message":"failed"}`, + } + + enqueueOpsErrorLog(ops, entry) + job := <-opsErrorLogQueue + require.LessOrEqual(t, len(job.entry.ErrorBody), service.OpsErrorLogQueueBodyMaxBytes) + require.NotContains(t, job.entry.ErrorBody, secret) + require.Equal(t, int64(1), OpsErrorLogSanitizedTotal()) +} + func TestOpsErrorLoggerMiddleware_DoesNotBreakOuterMiddlewares(t *testing.T) { gin.SetMode(gin.TestMode) @@ -149,6 +251,42 @@ func setupOpsErrorLogTestQueue(t *testing.T, size int) { opsErrorLogMu.Unlock() } +func TestOpsErrorLoggerMiddleware_HardSkipsIngressRejection(t *testing.T) { + setupOpsErrorLogTestQueue(t, 4) + gin.SetMode(gin.TestMode) + + settings := &ingressRejectSettingRepo{} + repo := &ingressRejectOpsRepo{} + ops := service.NewOpsService(repo, settings, nil, nil, nil, nil, nil, nil, nil, nil, nil) + // Construction may read unrelated runtime settings; only request-path reads matter here. + settings.getValueCalls = 0 + + router := gin.New() + router.Use(OpsErrorLoggerMiddleware(ops)) + router.GET("/v1/messages", func(c *gin.Context) { + middleware2.MarkIngressRejected(c, middleware2.IngressRejectInvalidAPIKey) + c.JSON(http.StatusUnauthorized, gin.H{"code": "INVALID_API_KEY", "message": "Invalid API key"}) + }) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/v1/messages", nil) + router.ServeHTTP(w, req) + + require.Equal(t, http.StatusUnauthorized, w.Code) + require.JSONEq(t, `{"code":"INVALID_API_KEY","message":"Invalid API key"}`, w.Body.String()) + require.Zero(t, settings.getValueCalls, "ingress rejection must bypass monitoring settings reads") + require.Zero(t, repo.insertCalls, "ingress rejection must bypass inserts") + require.Zero(t, OpsErrorLogEnqueuedTotal(), "ingress rejection must not enter the error queue") +} + +func TestNormalizeOpsPersistentUserAgentBoundsAndPreservesUTF8(t *testing.T) { + value := strings.Repeat("a", opsErrorLogMaxUserAgentBytes-1) + "你" + strings.Repeat("b", 32) + got := normalizeOpsPersistentUserAgent(" " + value + " ") + require.LessOrEqual(t, len(got), opsErrorLogMaxUserAgentBytes) + require.True(t, utf8.ValidString(got)) + require.NotContains(t, got, "b") +} + // 就地(in-band) SSE 错误挂在已固化的 HTTP 200 流上:wire 状态码为 200, // 常规 status>=400 采集路径不会触发。logOpsStreamError 必须据 MarkOpsStreamError // 补记一条错误日志,且用 IntendedStatus(429) 分级、StatusCode 仍记 wire 的 200。 @@ -182,6 +320,36 @@ func TestLogOpsStreamError_RecordsInBandConcurrencyLimit(t *testing.T) { require.Equal(t, "Concurrency limit exceeded for account, please retry later", job.entry.ErrorMessage) } +func TestLogOpsStreamError_UpstreamFailureCountsTowardsSLA(t *testing.T) { + setupOpsErrorLogTestQueue(t, 4) + + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + c.Set(opsModelKey, "gpt-5.6-sol") + + service.MarkOpsStreamFailure( + c, + "upstream_error", + service.OpenAIUpstreamHTTP2StreamErrorCode, + "Upstream HTTP/2 stream failed", + http.StatusBadGateway, + ) + + ops := service.NewOpsService(nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil) + logOpsStreamError(c, ops, http.StatusOK) + + job := <-opsErrorLogQueue + require.NotNil(t, job.entry) + require.Equal(t, http.StatusBadGateway, job.entry.StatusCode) + require.Equal(t, "upstream_error", job.entry.ErrorType) + require.Equal(t, "upstream", job.entry.ErrorPhase) + require.Equal(t, "provider", job.entry.ErrorOwner) + require.False(t, job.entry.IsBusinessLimited) + require.Contains(t, job.entry.ErrorBody, service.OpenAIUpstreamHTTP2StreamErrorCode) +} + // 未标记流内错误时 logOpsStreamError 必须是 no-op(不误记正常的 200 流)。 func TestLogOpsStreamError_NoopWhenNotMarked(t *testing.T) { setupOpsErrorLogTestQueue(t, 4) diff --git a/backend/internal/handler/ops_ingress_reject_capture_test.go b/backend/internal/handler/ops_ingress_reject_capture_test.go new file mode 100644 index 000000000..6848bcfd9 --- /dev/null +++ b/backend/internal/handler/ops_ingress_reject_capture_test.go @@ -0,0 +1,25 @@ +package handler + +import ( + "net/http" + "net/http/httptest" + "testing" + + middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestOpsCaptureWriterDoesNotCopyIngressRejectBody(t *testing.T) { + gin.SetMode(gin.TestMode) + context, _ := gin.CreateTestContext(httptest.NewRecorder()) + writer := acquireOpsCaptureWriter(context.Writer) + defer releaseOpsCaptureWriter(writer) + writer.ctx = context + context.Writer = writer + middleware2.MarkIngressRejected(context, middleware2.IngressRejectInvalidAPIKey) + context.Status(http.StatusUnauthorized) + _, err := context.Writer.WriteString(`{"code":"INVALID_API_KEY","message":"Invalid API key"}`) + require.NoError(t, err) + require.Zero(t, writer.buf.Len()) +} diff --git a/backend/internal/handler/wire.go b/backend/internal/handler/wire.go index 1f7186413..54d82a69e 100644 --- a/backend/internal/handler/wire.go +++ b/backend/internal/handler/wire.go @@ -120,12 +120,14 @@ func ProvideOpenAIGatewayHandler( errorPassthroughService *service.ErrorPassthroughService, contentModerationService *service.ContentModerationService, opsService *service.OpsService, + grokQuotaService *service.GrokQuotaService, cfg *config.Config, coordinator *securityaudit.Coordinator, ) *OpenAIGatewayHandler { h := NewOpenAIGatewayHandler(gatewayService, concurrencyService, billingCacheService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, opsService, cfg) h.securityAuditCoordinator = coordinator + h.grokMediaEligibilityProber = grokQuotaService return h } diff --git a/backend/internal/pkg/apicompat/anthropic_to_responses_response.go b/backend/internal/pkg/apicompat/anthropic_to_responses_response.go index 9ca20e1cd..b22215836 100644 --- a/backend/internal/pkg/apicompat/anthropic_to_responses_response.go +++ b/backend/internal/pkg/apicompat/anthropic_to_responses_response.go @@ -152,11 +152,25 @@ type AnthropicEventToResponsesState struct { // For message output: accumulate text parts ContentIndex int + // TextAccum accumulates the current text part so that output_text.done and + // content_part.done can carry the full text (deltas carry increments only). + TextAccum string // For function_call: track per-output info CurrentCallID string CurrentName string + // Content of the currently open item, folded into Outputs when it closes. + CurrentContent []ResponsesContentPart // message + CurrentArgs string // function_call + CurrentSummary string // reasoning + + // Outputs accumulates every closed output item so that response.completed + // can carry the full output list. The OpenAI SDK's get_final_response() + // parses the terminal event's response directly; without this, clients see + // an empty output_text. + Outputs []ResponsesOutput + // Usage from message_start / message_delta. InputTokens here follows // Anthropic semantics (excludes cached tokens); they are added back when // emitting the OpenAI Responses usage. @@ -293,6 +307,22 @@ func anthToResHandleContentBlockStart(evt *AnthropicStreamEvent, state *Anthropi })) } + // response.content_part.added must precede the output_text.delta events + // for that part. The message item is added with content: [], and the + // OpenAI SDK's accumulating stream helper (client.responses.stream) only + // appends a content part when it sees content_part.added. Without it the + // following output_text.delta indexes output.content[content_index] and + // raises IndexError. Raw event iteration + // (responses.create(stream=True)) does not accumulate, which is why this + // went unnoticed. + 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)...) @@ -327,6 +357,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, @@ -338,6 +369,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, @@ -349,6 +381,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, @@ -393,12 +426,24 @@ 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) + // Text block is done: emit output_text.done then content_part.done (the + // order OpenAI uses), both carrying the part's full text. The message + // item itself stays open since more blocks may follow. + text := state.TextAccum + state.TextAccum = "" + state.CurrentContent = append(state.CurrentContent, ResponsesContentPart{Type: "output_text", Text: text}) return []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}, }), } } @@ -454,24 +499,48 @@ func closeCurrentResponsesItem(state *AnthropicEventToResponsesState) []Response return nil } - itemType := state.CurrentItemType - itemID := state.CurrentItemID + // Assemble the full item: both output_item.done and response.completed must + // carry its content. Emitting only {type,id,status} makes SDK-side + // accumulation produce an empty output. + 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 + args := state.CurrentArgs + if args == "" { + args = "{}" + } + item.Arguments = args + 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, })} } @@ -519,6 +588,14 @@ func makeResponsesCompletedEvent( eventType = "response.incomplete" } + // Carry the output items accumulated over the stream. The SDK's + // get_final_response() reads them straight from the terminal event, so an + // empty list leaves clients with an empty result. + outputs := state.Outputs + if outputs == nil { + outputs = []ResponsesOutput{} + } + return ResponsesStreamEvent{ Type: eventType, SequenceNumber: seq, @@ -527,7 +604,7 @@ func makeResponsesCompletedEvent( Object: "response", Model: state.Model, Status: status, - Output: []ResponsesOutput{}, + 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..c26e28982 --- /dev/null +++ b/backend/internal/pkg/apicompat/anthropic_to_responses_stream_test.go @@ -0,0 +1,187 @@ +package apicompat + +import "testing" + +// TestAnthropicEventToResponses_TextEmitsContentPart pins that a message text +// stream emits response.content_part.added, and that it precedes the first +// output_text.delta for that part. +// +// Why: the OpenAI SDK's accumulating stream helper (client.responses.stream) +// only appends a content part to the message item when it sees +// content_part.added. The item is added with content: [], so a missing event +// makes the following output_text.delta index output.content[content_index] and +// raise IndexError. Raw event iteration does not accumulate, so a regression +// here is easy to miss. +func TestAnthropicEventToResponses_TextEmitsContentPart(t *testing.T) { + state := NewAnthropicEventToResponsesState() + state.Model = "claude-sonnet-4-5" + + var types []string + feed := func(evt *AnthropicStreamEvent) { + for _, out := range AnthropicEventToResponsesEvents(evt, state) { + types = append(types, out.Type) + } + } + + idx := 0 + feed(&AnthropicStreamEvent{Type: "message_start", Message: &AnthropicResponse{ID: "msg_1", Model: "claude-sonnet-4-5"}}) + feed(&AnthropicStreamEvent{Type: "content_block_start", Index: &idx, ContentBlock: &AnthropicContentBlock{Type: "text"}}) + feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &idx, Delta: &AnthropicDelta{Type: "text_delta", Text: "Hel"}}) + feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &idx, Delta: &AnthropicDelta{Type: "text_delta", Text: "lo"}}) + feed(&AnthropicStreamEvent{Type: "content_block_stop", Index: &idx}) + feed(&AnthropicStreamEvent{Type: "message_stop"}) + + posOf := func(target string) int { + for i, ty := range types { + if ty == target { + return i + } + } + return -1 + } + + partAdded := posOf("response.content_part.added") + firstDelta := posOf("response.output_text.delta") + + if partAdded < 0 { + t.Fatalf("response.content_part.added was not emitted; got %v", types) + } + if firstDelta < 0 { + t.Fatalf("response.output_text.delta was not emitted; got %v", types) + } + if partAdded > firstDelta { + t.Errorf("content_part.added must precede the first output_text.delta; got %v", types) + } + if posOf("response.content_part.done") < 0 { + t.Errorf("response.content_part.done was not emitted; got %v", types) + } +} + +// TestAnthropicEventToResponses_DoneEventsCarryFullText pins that done events +// carry the part's full text (deltas carry increments only). +func TestAnthropicEventToResponses_DoneEventsCarryFullText(t *testing.T) { + state := NewAnthropicEventToResponsesState() + state.Model = "claude-sonnet-4-5" + + var events []ResponsesStreamEvent + feed := func(evt *AnthropicStreamEvent) { + events = append(events, AnthropicEventToResponsesEvents(evt, state)...) + } + + idx := 0 + feed(&AnthropicStreamEvent{Type: "message_start", Message: &AnthropicResponse{ID: "msg_1"}}) + feed(&AnthropicStreamEvent{Type: "content_block_start", Index: &idx, ContentBlock: &AnthropicContentBlock{Type: "text"}}) + feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &idx, Delta: &AnthropicDelta{Type: "text_delta", Text: "Hello "}}) + feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &idx, Delta: &AnthropicDelta{Type: "text_delta", Text: "world"}}) + feed(&AnthropicStreamEvent{Type: "content_block_stop", Index: &idx}) + + const want = "Hello world" + var sawTextDone, sawPartDone bool + for _, e := range events { + switch e.Type { + case "response.output_text.done": + sawTextDone = true + if e.Text != want { + t.Errorf("output_text.done text = %q, want %q", e.Text, want) + } + case "response.content_part.done": + sawPartDone = true + if e.Part == nil || e.Part.Text != want { + t.Errorf("content_part.done part = %+v, want text %q", e.Part, want) + } + } + } + if !sawTextDone || !sawPartDone { + t.Errorf("missing done events: output_text.done=%v content_part.done=%v", sawTextDone, sawPartDone) + } +} + +// TestAnthropicEventToResponses_CompletedCarriesOutput pins that +// response.completed carries the full output list. The SDK's +// get_final_response() and tracing integrations parse the terminal event's +// response directly; an empty output leaves them with nothing (the text still +// renders from deltas, which is why this is invisible when only watching the +// stream). +func TestAnthropicEventToResponses_CompletedCarriesOutput(t *testing.T) { + state := NewAnthropicEventToResponsesState() + state.Model = "claude-sonnet-4-5" + + var events []ResponsesStreamEvent + feed := func(evt *AnthropicStreamEvent) { + events = append(events, AnthropicEventToResponsesEvents(evt, state)...) + } + + idx := 0 + feed(&AnthropicStreamEvent{Type: "message_start", Message: &AnthropicResponse{ID: "msg_1"}}) + feed(&AnthropicStreamEvent{Type: "content_block_start", Index: &idx, ContentBlock: &AnthropicContentBlock{Type: "text"}}) + feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &idx, Delta: &AnthropicDelta{Type: "text_delta", Text: "4826"}}) + feed(&AnthropicStreamEvent{Type: "content_block_stop", Index: &idx}) + feed(&AnthropicStreamEvent{Type: "message_stop"}) + + var completed *ResponsesStreamEvent + for i := range events { + if events[i].Type == "response.completed" { + completed = &events[i] + } + } + if completed == nil || completed.Response == nil { + t.Fatalf("response.completed was not emitted") + } + if len(completed.Response.Output) == 0 { + t.Fatalf("response.completed carries an empty output; clients would see no result") + } + msg := completed.Response.Output[0] + if msg.Type != "message" || len(msg.Content) == 0 { + t.Fatalf("output[0] = %+v, want a message with content", msg) + } + if msg.Content[0].Text != "4826" { + t.Errorf("output[0].content[0].text = %q, want %q", msg.Content[0].Text, "4826") + } +} + +// TestAnthropicEventToResponses_ToolCallCompletedCarriesArguments pins that a +// function call's accumulated arguments survive into output_item.done and +// response.completed. +func TestAnthropicEventToResponses_ToolCallCompletedCarriesArguments(t *testing.T) { + state := NewAnthropicEventToResponsesState() + state.Model = "claude-sonnet-4-5" + + var events []ResponsesStreamEvent + feed := func(evt *AnthropicStreamEvent) { + events = append(events, AnthropicEventToResponsesEvents(evt, state)...) + } + + idx := 0 + feed(&AnthropicStreamEvent{Type: "message_start", Message: &AnthropicResponse{ID: "msg_1"}}) + feed(&AnthropicStreamEvent{Type: "content_block_start", Index: &idx, ContentBlock: &AnthropicContentBlock{ + Type: "tool_use", ID: "toolu_1", Name: "get_weather", + }}) + feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &idx, Delta: &AnthropicDelta{ + Type: "input_json_delta", PartialJSON: `{"city":`, + }}) + feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &idx, Delta: &AnthropicDelta{ + Type: "input_json_delta", PartialJSON: `"SH"}`, + }}) + feed(&AnthropicStreamEvent{Type: "content_block_stop", Index: &idx}) + feed(&AnthropicStreamEvent{Type: "message_stop"}) + + var completed *ResponsesStreamEvent + for i := range events { + if events[i].Type == "response.completed" { + completed = &events[i] + } + } + if completed == nil || completed.Response == nil || len(completed.Response.Output) == 0 { + t.Fatalf("response.completed carries no output") + } + fc := completed.Response.Output[0] + if fc.Type != "function_call" { + t.Fatalf("output[0].type = %q, want function_call", fc.Type) + } + if fc.Arguments != `{"city":"SH"}` { + t.Errorf("arguments = %q, want %q", fc.Arguments, `{"city":"SH"}`) + } + if fc.Name != "get_weather" { + t.Errorf("name = %q, want get_weather", fc.Name) + } +} diff --git a/backend/internal/pkg/ip/ip.go b/backend/internal/pkg/ip/ip.go index 2702b1353..b9c777449 100644 --- a/backend/internal/pkg/ip/ip.go +++ b/backend/internal/pkg/ip/ip.go @@ -8,40 +8,10 @@ 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 resolves a client address only through Gin's configured trusted +// proxy chain. Forwarding headers from a direct or untrusted peer are ignored. 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 GetTrustedClientIP(c) } // GetTrustedClientIP 从 Gin 的可信代理解析链提取客户端 IP。 @@ -54,14 +24,10 @@ func GetTrustedClientIP(c *gin.Context) string { return normalizeIP(c.ClientIP()) } -// GetSecurityClientIP 返回安全敏感场景(API Key IP 限制、审计日志、会话 IP/UA 绑定) -// 使用的客户端 IP。trustForwarded 对应系统设置「信任反代传递的客户端 IP」: -// 开启时信任反代转发头(CF-Connecting-IP / X-Real-IP / X-Forwarded-For), -// 关闭时走 Gin trusted_proxies 解析链。 -func GetSecurityClientIP(c *gin.Context, trustForwarded bool) string { - if trustForwarded { - return GetClientIP(c) - } +// GetSecurityClientIP returns the address resolved through Gin's configured +// trusted-proxy chain. The legacy toggle is retained for configuration/API +// compatibility, but never makes raw forwarding headers trustworthy by itself. +func GetSecurityClientIP(c *gin.Context, _ bool) string { return GetTrustedClientIP(c) } @@ -75,9 +41,6 @@ func normalizeIP(ip string) string { return ip } -// privateNets 预编译私有 IP CIDR 块,避免每次调用 isPrivateIP 时重复解析 -var privateNets []*net.IPNet - // CompiledIPRules 表示预编译的 IP 匹配规则。 // PatternCount 记录原始规则数量,用于保留“规则存在但全无效”时的行为语义。 type CompiledIPRules struct { @@ -86,23 +49,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 { @@ -150,20 +96,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 7e579bf7a..2432abe7a 100644 --- a/backend/internal/pkg/ip/ip_test.go +++ b/backend/internal/pkg/ip/ip_test.go @@ -10,48 +10,6 @@ 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) { gin.SetMode(gin.TestMode) @@ -95,7 +53,7 @@ func TestCheckIPRestrictionWithCompiledRules_InvalidWhitelistStillDenies(t *test require.Equal(t, "access denied", reason) } -func TestGetSecurityClientIPHonorsTrustToggle(t *testing.T) { +func TestGetSecurityClientIPNeverTrustsHeadersFromUntrustedPeer(t *testing.T) { gin.SetMode(gin.TestMode) for _, tc := range []struct { @@ -103,8 +61,8 @@ func TestGetSecurityClientIPHonorsTrustToggle(t *testing.T) { trustForwarded bool want string }{ - {name: "trust disabled uses trusted proxy chain", trustForwarded: false, want: "9.9.9.9"}, - {name: "trust enabled uses forwarded header", trustForwarded: true, want: "1.2.3.4"}, + {name: "legacy toggle disabled", trustForwarded: false, want: "9.9.9.9"}, + {name: "legacy toggle enabled", trustForwarded: true, want: "9.9.9.9"}, } { t.Run(tc.name, func(t *testing.T) { r := gin.New() @@ -124,3 +82,18 @@ func TestGetSecurityClientIPHonorsTrustToggle(t *testing.T) { }) } } + +func TestGetSecurityClientIPUsesConfiguredTrustedProxy(t *testing.T) { + gin.SetMode(gin.TestMode) + r := gin.New() + require.NoError(t, r.SetTrustedProxies([]string{"9.9.9.9"})) + r.GET("/t", func(c *gin.Context) { c.String(200, GetSecurityClientIP(c, true)) }) + + w := httptest.NewRecorder() + req := httptest.NewRequest("GET", "/t", nil) + req.RemoteAddr = "9.9.9.9:12345" + req.Header.Set("X-Forwarded-For", "1.2.3.4") + r.ServeHTTP(w, req) + + require.Equal(t, "1.2.3.4", w.Body.String()) +} diff --git a/backend/internal/pkg/logger/logger.go b/backend/internal/pkg/logger/logger.go index 3fca706ec..de08dab35 100644 --- a/backend/internal/pkg/logger/logger.go +++ b/backend/internal/pkg/logger/logger.go @@ -26,6 +26,10 @@ const ( LevelWarn = zapcore.WarnLevel LevelError = zapcore.ErrorLevel LevelFatal = zapcore.FatalLevel + + // OpsSystemLogSkipField keeps an event in the standard logger while + // preventing the database-backed Ops system-log sink from indexing it. + OpsSystemLogSkipField = "ops_system_log_skip" ) type Sink interface { diff --git a/backend/internal/repository/account_repo.go b/backend/internal/repository/account_repo.go index f014155bb..87f91777b 100644 --- a/backend/internal/repository/account_repo.go +++ b/backend/internal/repository/account_repo.go @@ -893,6 +893,20 @@ func (r *accountRepository) ListOpsAccountsForStats(ctx context.Context, platfor func accountListOrder(params pagination.PaginationParams) []func(*entsql.Selector) { sortBy := strings.ToLower(strings.TrimSpace(params.SortBy)) sortOrder := params.NormalizedSortOrder(pagination.SortOrderAsc) + if sortBy == "upstream_billing_rate" { + direction := "ASC" + tieOrder := entsql.Asc + if sortOrder == pagination.SortOrderDesc { + direction = "DESC" + tieOrder = entsql.Desc + } + return []func(*entsql.Selector){func(s *entsql.Selector) { + extra := s.C(dbaccount.FieldExtra) + expression := upstreamBillingRateSortExpression(extra) + s.OrderExpr(entsql.Expr(expression + " " + direction + " NULLS LAST")) + s.OrderBy(tieOrder(s.C(dbaccount.FieldID))) + }} + } field := dbaccount.FieldName defaultOrder := true @@ -934,6 +948,40 @@ func accountListOrder(params pagination.PaginationParams) []func(*entsql.Selecto return []func(*entsql.Selector){dbent.Asc(field), dbent.Asc(dbaccount.FieldID)} } +func upstreamBillingRateSortExpression(extra string) string { + status := extra + " #>> '{upstream_billing_probe,status}'" + effectiveJSON := extra + " #> '{upstream_billing_probe,data,effective_rate_multiplier}'" + effective := extra + " #>> '{upstream_billing_probe,data,effective_rate_multiplier}'" + resolvedJSON := extra + " #> '{upstream_billing_probe,data,resolved_rate_multiplier}'" + resolved := extra + " #>> '{upstream_billing_probe,data,resolved_rate_multiplier}'" + peakEnabledJSON := extra + " #> '{upstream_billing_probe,data,peak_rate_enabled}'" + peakEnabled := extra + " #>> '{upstream_billing_probe,data,peak_rate_enabled}'" + peakStart := extra + " #>> '{upstream_billing_probe,data,peak_start}'" + peakEnd := extra + " #>> '{upstream_billing_probe,data,peak_end}'" + peakMultiplierJSON := extra + " #> '{upstream_billing_probe,data,peak_rate_multiplier}'" + peakMultiplier := extra + " #>> '{upstream_billing_probe,data,peak_rate_multiplier}'" + peakMultiplierValue := "(CASE WHEN jsonb_typeof(" + peakMultiplierJSON + ") = 'number' THEN (" + peakMultiplier + ")::numeric END)" + billingScope := extra + " #>> '{upstream_billing_probe,data,billing_scope}'" + timezone := extra + " #>> '{upstream_billing_probe,data,timezone}'" + validClock := "'^([01][0-9]|2[0-3]):[0-5][0-9]$'" + startMinute := "(CASE WHEN " + peakStart + " ~ " + validClock + " THEN split_part(" + peakStart + ", ':', 1)::numeric * 60 + split_part(" + peakStart + ", ':', 2)::numeric END)" + endMinute := "(CASE WHEN " + peakEnd + " ~ " + validClock + " THEN split_part(" + peakEnd + ", ':', 1)::numeric * 60 + split_part(" + peakEnd + ", ':', 2)::numeric END)" + localMinute := "(EXTRACT(HOUR FROM (CURRENT_TIMESTAMP AT TIME ZONE (" + timezone + "))) * 60 + EXTRACT(MINUTE FROM (CURRENT_TIMESTAMP AT TIME ZONE (" + timezone + "))))" + validPeakWindow := peakStart + " ~ " + validClock + " AND " + + peakEnd + " ~ " + validClock + " AND " + + startMinute + " < " + endMinute + validPeakConfig := validPeakWindow + " AND " + peakMultiplierValue + " >= 0 AND " + + "EXISTS (SELECT 1 FROM pg_timezone_names WHERE name = " + timezone + ")" + dynamicRate := "CASE WHEN " + peakEnabled + " = 'false' THEN (" + resolved + ")::numeric WHEN " + peakEnabled + " = 'true' AND " + validPeakConfig + + " THEN (" + resolved + ")::numeric * CASE WHEN " + localMinute + " >= " + startMinute + " AND " + localMinute + " < " + endMinute + + " THEN " + peakMultiplierValue + " ELSE 1 END ELSE NULL END" + legacySnapshot := "jsonb_typeof(" + resolvedJSON + ") IS NULL AND jsonb_typeof(" + peakEnabledJSON + ") IS NULL" + + return "CASE WHEN " + status + " IN ('ok', 'failed') AND (jsonb_typeof(" + resolvedJSON + ") = 'number' OR jsonb_typeof(" + effectiveJSON + ") = 'number') THEN CASE WHEN jsonb_typeof(" + + resolvedJSON + ") = 'number' AND jsonb_typeof(" + peakEnabledJSON + ") = 'boolean' THEN CASE WHEN " + billingScope + " = 'token' THEN " + dynamicRate + " ELSE NULL END WHEN " + legacySnapshot + + " AND jsonb_typeof(" + effectiveJSON + ") = 'number' THEN (" + effective + ")::numeric END END" +} + func (r *accountRepository) ListByGroup(ctx context.Context, groupID int64) ([]service.Account, error) { accounts, err := r.queryAccountsByGroup(ctx, groupID, accountGroupQueryOptions{ status: service.StatusActive, @@ -1877,6 +1925,46 @@ func (r *accountRepository) ListSchedulableByGroupIDAndPlatforms(ctx context.Con }) } +// ListModelAvailabilityCandidates returns the persistently configured account +// pool used to decide whether a model is supported. Unlike scheduling queries, +// it intentionally ignores transient runtime state (rate limits, overload, +// temporary unschedulability, and expiry windows). +func (r *accountRepository) ListModelAvailabilityCandidates( + ctx context.Context, + groupID *int64, + platforms []string, + includeGrouped bool, +) ([]service.Account, error) { + if len(platforms) == 0 { + return []service.Account{}, nil + } + if groupID != nil { + return r.queryAccountsByGroup(ctx, *groupID, accountGroupQueryOptions{ + status: service.StatusActive, + schedulable: true, + ignoreTransientState: true, + platforms: platforms, + }) + } + + preds := []dbpredicate.Account{ + dbaccount.StatusEQ(service.StatusActive), + dbaccount.SchedulableEQ(true), + dbaccount.PlatformIn(platforms...), + } + if !includeGrouped { + preds = append(preds, dbaccount.Not(dbaccount.HasAccountGroups())) + } + accounts, err := r.client.Account.Query(). + Where(preds...). + Order(dbent.Asc(dbaccount.FieldPriority)). + All(ctx) + if err != nil { + return nil, err + } + return r.accountsToService(ctx, accounts) +} + func (r *accountRepository) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error { now := time.Now() _, err := r.client.Account.Update(). @@ -2575,6 +2663,12 @@ func (r *accountRepository) BulkUpdate(ctx context.Context, ids []int64, updates args = append(args, *updates.Schedulable) idx++ } + if updates.ProbeEnabled != nil { + if updates.Extra == nil { + updates.Extra = make(map[string]any) + } + updates.Extra[service.UpstreamBillingProbeEnabledExtraKey] = *updates.ProbeEnabled + } // JSONB 需要合并而非覆盖,使用 raw SQL 保持旧行为。 if len(updates.Credentials) > 0 { payload, err := json.Marshal(updates.Credentials) @@ -2605,8 +2699,14 @@ func (r *accountRepository) BulkUpdate(ctx context.Context, ids []int64, updates setClauses = append(setClauses, "updated_at = NOW()") - query := "UPDATE accounts SET " + joinClauses(setClauses, ", ") + " WHERE id = ANY($" + itoa(idx) + ") AND deleted_at IS NULL" + whereClause := " WHERE id = ANY($" + itoa(idx) + ") AND deleted_at IS NULL" args = append(args, pq.Array(ids)) + idx++ + if updates.ProbeEnabled != nil { + whereClause += " AND platform = $" + itoa(idx) + " AND type = $" + itoa(idx+1) + args = append(args, service.PlatformOpenAI, service.AccountTypeAPIKey) + } + query := "UPDATE accounts SET " + joinClauses(setClauses, ", ") + whereClause baseCtx := ctx contextTx := dbent.TxFromContext(ctx) @@ -2635,6 +2735,20 @@ func (r *accountRepository) BulkUpdate(ctx context.Context, ids []int64, updates if err != nil { return 0, err } + if updates.ProbeEnabled != nil { + expectedRows := int64(0) + seenIDs := make(map[int64]struct{}, len(ids)) + for _, id := range ids { + if _, seen := seenIDs[id]; seen { + continue + } + seenIDs[id] = struct{}{} + expectedRows++ + } + if rows != expectedRows { + return 0, service.ErrUpstreamBillingProbeAccountInvalid + } + } if rows > 0 { payload := map[string]any{"account_ids": ids} if err := enqueueSchedulerOutbox(ctx, exec, service.SchedulerOutboxEventAccountBulkChanged, nil, nil, payload); err != nil { @@ -2662,9 +2776,10 @@ func (r *accountRepository) BulkUpdate(ctx context.Context, ids []int64, updates } type accountGroupQueryOptions struct { - status string - schedulable bool - platforms []string // 允许的多个平台,空切片表示不进行平台过滤 + status string + schedulable bool + ignoreTransientState bool + platforms []string // 允许的多个平台,空切片表示不进行平台过滤 } func (r *accountRepository) queryAccountsByGroup(ctx context.Context, groupID int64, opts accountGroupQueryOptions) ([]service.Account, error) { @@ -2681,14 +2796,16 @@ func (r *accountRepository) queryAccountsByGroup(ctx context.Context, groupID in preds = append(preds, dbaccount.PlatformIn(opts.platforms...)) } if opts.schedulable { - now := time.Now() - preds = append(preds, - dbaccount.SchedulableEQ(true), - tempUnschedulablePredicate(), - notExpiredPredicate(now), - dbaccount.Or(dbaccount.OverloadUntilIsNil(), dbaccount.OverloadUntilLTE(now)), - dbaccount.Or(dbaccount.RateLimitResetAtIsNil(), dbaccount.RateLimitResetAtLTE(now)), - ) + preds = append(preds, dbaccount.SchedulableEQ(true)) + if !opts.ignoreTransientState { + now := time.Now() + preds = append(preds, + tempUnschedulablePredicate(), + notExpiredPredicate(now), + dbaccount.Or(dbaccount.OverloadUntilIsNil(), dbaccount.OverloadUntilLTE(now)), + dbaccount.Or(dbaccount.RateLimitResetAtIsNil(), dbaccount.RateLimitResetAtLTE(now)), + ) + } } if len(preds) > 0 { diff --git a/backend/internal/repository/account_repo_model_availability_test.go b/backend/internal/repository/account_repo_model_availability_test.go new file mode 100644 index 000000000..917e73b0c --- /dev/null +++ b/backend/internal/repository/account_repo_model_availability_test.go @@ -0,0 +1,58 @@ +package repository + +import ( + "context" + "strings" + "testing" + + "github.com/DATA-DOG/go-sqlmock" + dbent "github.com/Wei-Shaw/sub2api/ent" + _ "github.com/Wei-Shaw/sub2api/ent/runtime" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/stretchr/testify/require" + + "entgo.io/ent/dialect" + entsql "entgo.io/ent/dialect/sql" +) + +func TestListModelAvailabilityCandidates_GroupQueryIgnoresTransientState(t *testing.T) { + var capturedSQL string + db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(captureEntQueryMatcher{actual: &capturedSQL})) + require.NoError(t, err) + t.Cleanup(func() { _ = db.Close() }) + + driver := entsql.OpenDB(dialect.Postgres, db) + client := dbent.NewClient(dbent.Driver(driver)) + t.Cleanup(func() { _ = client.Close() }) + repo := newAccountRepositoryWithSQL(client, db, nil) + + mock.ExpectQuery("model availability candidates"). + WillReturnRows(sqlmock.NewRows([]string{"id"})) + groupID := int64(42) + accounts, err := repo.ListModelAvailabilityCandidates( + context.Background(), + &groupID, + []string{service.PlatformAnthropic}, + false, + ) + require.NoError(t, err) + require.Empty(t, accounts) + require.NoError(t, mock.ExpectationsWereMet()) + + normalized := normalizeSQLWhitespace(capturedSQL) + _, whereClause, found := strings.Cut(normalized, " WHERE ") + require.True(t, found, "expected WHERE clause in query: %s", normalized) + whereClause, _, _ = strings.Cut(whereClause, " ORDER BY ") + for _, configuredPredicate := range []string{"group_id", "status", "schedulable", "platform"} { + require.Contains(t, whereClause, configuredPredicate) + } + for _, transientPredicate := range []string{ + "rate_limit_reset_at", + "overload_until", + "temp_unschedulable_until", + "expires_at", + "auto_pause_on_expired", + } { + require.NotContains(t, whereClause, transientPredicate, "configured-state diagnosis must not filter transient predicate %q", transientPredicate) + } +} diff --git a/backend/internal/repository/account_repo_sort_integration_test.go b/backend/internal/repository/account_repo_sort_integration_test.go index 098dde7bb..19c974f37 100644 --- a/backend/internal/repository/account_repo_sort_integration_test.go +++ b/backend/internal/repository/account_repo_sort_integration_test.go @@ -3,6 +3,9 @@ package repository import ( + "fmt" + "time" + "github.com/Wei-Shaw/sub2api/internal/pkg/pagination" "github.com/Wei-Shaw/sub2api/internal/service" ) @@ -33,3 +36,105 @@ func (s *AccountRepoSuite) TestListWithFilters_SortByPriorityDesc() { s.Require().Equal("high-priority", accounts[0].Name) s.Require().Equal("low-priority", accounts[1].Name) } + +func (s *AccountRepoSuite) TestListWithFilters_SortByUpstreamBillingRateWithMissingLast() { + makeAccount := func(name, status string, rate any) { + extra := map[string]any{} + if rate != nil { + extra[service.UpstreamBillingProbeExtraKey] = map[string]any{ + "status": status, + "data": map[string]any{"effective_rate_multiplier": rate}, + } + } + mustCreateAccount(s.T(), s.client, &service.Account{Name: name, Extra: extra}) + } + makeAccount("high-rate", service.UpstreamBillingProbeStatusOK, 0.8) + makeAccount("low-rate", service.UpstreamBillingProbeStatusOK, 0.03) + makeAccount("missing-rate", "", nil) + makeAccount("unsupported-with-retained-rate", service.UpstreamBillingProbeStatusUnsupported, 0.01) + + for _, tc := range []struct { + order string + want []string + }{ + {order: "asc", want: []string{"low-rate", "high-rate", "missing-rate", "unsupported-with-retained-rate"}}, + {order: "desc", want: []string{"high-rate", "low-rate", "unsupported-with-retained-rate", "missing-rate"}}, + } { + accounts, _, err := s.repo.ListWithFilters(s.ctx, pagination.PaginationParams{ + Page: 1, PageSize: 10, SortBy: "upstream_billing_rate", SortOrder: tc.order, + }, "", "", "", "", 0, "") + s.Require().NoError(err) + s.Require().Len(accounts, 4) + for i, name := range tc.want { + s.Require().Equal(name, accounts[i].Name) + } + } +} + +func (s *AccountRepoSuite) TestListWithFilters_SortByCurrentUpstreamBillingRateDuringPeak() { + now := time.Now() + locations := []string{"UTC", "Asia/Shanghai", "America/New_York", "Europe/London"} + var timezone string + var minute int + for _, name := range locations { + location, err := time.LoadLocation(name) + s.Require().NoError(err) + local := now.In(location) + candidate := local.Hour()*60 + local.Minute() + if candidate >= 2 && candidate <= 1436 { + timezone = name + minute = candidate + break + } + } + s.Require().NotEmpty(timezone) + + peakStart := fmt.Sprintf("%02d:%02d", (minute-2)/60, (minute-2)%60) + peakEnd := fmt.Sprintf("%02d:%02d", (minute+3)/60, (minute+3)%60) + mustCreateAccount(s.T(), s.client, &service.Account{ + Name: "current-peak-rate", + Extra: map[string]any{ + service.UpstreamBillingProbeExtraKey: map[string]any{ + "status": service.UpstreamBillingProbeStatusOK, + "data": map[string]any{ + "billing_scope": "token", + "resolved_rate_multiplier": 1.0, + "effective_rate_multiplier": 1.0, + "peak_rate_enabled": true, + "peak_start": peakStart, + "peak_end": peakEnd, + "peak_rate_multiplier": 10.0, + "timezone": timezone, + }, + }, + }, + }) + mustCreateAccount(s.T(), s.client, &service.Account{ + Name: "current-off-peak-rate", + Extra: map[string]any{ + service.UpstreamBillingProbeExtraKey: map[string]any{ + "status": service.UpstreamBillingProbeStatusOK, + "data": map[string]any{ + "effective_rate_multiplier": 5.0, + }, + }, + }, + }) + + for _, tc := range []struct { + order string + want []string + }{ + {order: "asc", want: []string{"current-off-peak-rate", "current-peak-rate"}}, + {order: "desc", want: []string{"current-peak-rate", "current-off-peak-rate"}}, + } { + accounts, _, err := s.repo.ListWithFilters(s.ctx, pagination.PaginationParams{ + Page: 1, PageSize: 10, SortBy: "upstream_billing_rate", SortOrder: tc.order, + }, "", "", "", "", 0, "") + s.Require().NoError(err) + s.Require().Len(accounts, 2) + for i, name := range tc.want { + s.Require().Equal(name, accounts[i].Name) + } + } +} diff --git a/backend/internal/repository/account_repo_upstream_billing_probe_update_test.go b/backend/internal/repository/account_repo_upstream_billing_probe_update_test.go index 26fbb74dc..2a78f38c1 100644 --- a/backend/internal/repository/account_repo_upstream_billing_probe_update_test.go +++ b/backend/internal/repository/account_repo_upstream_billing_probe_update_test.go @@ -163,6 +163,46 @@ func TestBulkUpdateNilProbeRemovesKeyInsteadOfWritingJSONNull(t *testing.T) { require.Contains(t, normalizeSQLWhitespace(exec.execQueries[0]), "- 'upstream_billing_probe'") } +func TestBulkUpdateDisablingProbeRemovesSnapshot(t *testing.T) { + exec := &recordingSQLExecutor{result: rowsAffectedResult(1)} + repo := newAccountRepositoryWithSQL(nil, exec, nil) + + _, err := repo.BulkUpdate(context.Background(), []int64{27}, service.AccountBulkUpdate{ + Extra: map[string]any{service.UpstreamBillingProbeEnabledExtraKey: false}, + }) + + require.NoError(t, err) + require.NotEmpty(t, exec.execQueries) + require.Contains(t, normalizeSQLWhitespace(exec.execQueries[0]), "- 'upstream_billing_probe'") + payload, ok := exec.execArgs[0][0].([]byte) + require.True(t, ok) + require.Equal(t, `{"upstream_billing_probe_enabled":false}`, string(payload)) +} + +func TestBulkUpdateProbeEligibilityMismatchRollsBack(t *testing.T) { + db, mock, err := sqlmock.New() + require.NoError(t, err) + t.Cleanup(func() { _ = db.Close() }) + client := dbent.NewClient(dbent.Driver(entsql.OpenDB(dialect.Postgres, db))) + t.Cleanup(func() { _ = client.Close() }) + + enabled := true + mock.ExpectBegin() + mock.ExpectExec(`(?s)UPDATE accounts SET extra = .* WHERE id = ANY\(\$2\) AND deleted_at IS NULL AND platform = \$3 AND type = \$4`). + WithArgs(sqlmock.AnyArg(), `{27,28}`, service.PlatformOpenAI, service.AccountTypeAPIKey). + WillReturnResult(sqlmock.NewResult(0, 1)) + mock.ExpectRollback() + + repo := newAccountRepositoryWithSQL(client, db, nil) + rows, err := repo.BulkUpdate(context.Background(), []int64{27, 28}, service.AccountBulkUpdate{ + ProbeEnabled: &enabled, + }) + + require.ErrorIs(t, err, service.ErrUpstreamBillingProbeAccountInvalid) + require.Zero(t, rows) + require.NoError(t, mock.ExpectationsWereMet()) +} + func TestUpdateCredentialsAtomicallyClearsProbeForOpenAIAPIKeyIdentityChange(t *testing.T) { db, mock, err := sqlmock.New() require.NoError(t, err) diff --git a/backend/internal/repository/api_key_cache.go b/backend/internal/repository/api_key_cache.go index a1072057e..ecc6026f4 100644 --- a/backend/internal/repository/api_key_cache.go +++ b/backend/internal/repository/api_key_cache.go @@ -110,28 +110,25 @@ func (c *apiKeyCache) SubscribeAuthCacheInvalidation(ctx context.Context, handle return fmt.Errorf("subscribe to auth cache invalidation: %w", err) } - go func() { - defer func() { - if err := pubsub.Close(); err != nil { - log.Printf("Warning: failed to close auth cache invalidation pubsub: %v", err) - } - }() - - ch := pubsub.Channel() - for { - select { - case <-ctx.Done(): - return - case msg, ok := <-ch: - if !ok { - return - } - if msg != nil { - handler(msg.Payload) - } - } + defer func() { + if err := pubsub.Close(); err != nil { + log.Printf("Warning: failed to close auth cache invalidation pubsub: %v", err) } }() + service.NotifyAuthCacheSubscriptionReady(ctx) - return nil + ch := pubsub.Channel() + for { + select { + case <-ctx.Done(): + return ctx.Err() + case msg, ok := <-ch: + if !ok { + return errors.New("auth cache invalidation pubsub channel closed") + } + if msg != nil { + handler(msg.Payload) + } + } + } } diff --git a/backend/internal/repository/api_key_cache_subscriber_test.go b/backend/internal/repository/api_key_cache_subscriber_test.go new file mode 100644 index 000000000..2a49a9369 --- /dev/null +++ b/backend/internal/repository/api_key_cache_subscriber_test.go @@ -0,0 +1,49 @@ +package repository + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/alicebob/miniredis/v2" + "github.com/redis/go-redis/v9" + "github.com/stretchr/testify/require" +) + +func TestAPIKeyCacheSubscriber_BlocksUntilContextCancellation(t *testing.T) { + server := miniredis.RunT(t) + client := redis.NewClient(&redis.Options{Addr: server.Addr()}) + defer func() { _ = client.Close() }() + cache := NewAPIKeyCache(client) + ctx, cancel := context.WithCancel(context.Background()) + received := make(chan string, 1) + returned := make(chan error, 1) + go func() { + returned <- cache.SubscribeAuthCacheInvalidation(ctx, func(value string) { received <- value }) + }() + + var value string + require.Eventually(t, func() bool { + require.NoError(t, client.Publish(context.Background(), authCacheInvalidateChannel, "hash").Err()) + select { + case value = <-received: + return true + default: + return false + } + }, time.Second, 10*time.Millisecond) + require.Equal(t, "hash", value) + select { + case err := <-returned: + t.Fatalf("subscriber returned while connection was active: %v", err) + default: + } + cancel() + select { + case err := <-returned: + require.True(t, errors.Is(err, context.Canceled)) + case <-time.After(time.Second): + t.Fatal("subscriber did not stop after context cancellation") + } +} diff --git a/backend/internal/repository/api_key_repo.go b/backend/internal/repository/api_key_repo.go index ece9aae77..c926a59f3 100644 --- a/backend/internal/repository/api_key_repo.go +++ b/backend/internal/repository/api_key_repo.go @@ -326,16 +326,14 @@ func (r *apiKeyRepository) Delete(ctx context.Context, id int64) error { return nil } -// DeleteWithAudit 在同一事务内: -// 1. 把(明文 key、所有者、key 名称)写入 deleted_api_key_audits; -// 2. 软删除该 key(tombstone 覆盖 key 列以释放唯一约束)。 -// -// 保证"被删除的 key 一定能反查到所有者"。事务模式与 group_repo.DeleteCascade 一致。 +// DeleteWithAudit keeps the legacy method name for rolling-upgrade compatibility. +// It atomically tombstones and soft-deletes the key without retaining credential +// material. Tombstoning releases the unique key value for safe reuse. func (r *apiKeyRepository) DeleteWithAudit(ctx context.Context, id int64) error { tombstoneKey := fmt.Sprintf("__deleted__%d__%d", id, time.Now().UnixNano()) if existingTx := dbent.TxFromContext(ctx); existingTx != nil { - return r.deleteWithAudit(ctx, existingTx.Client(), id, tombstoneKey) + return r.deleteWithTombstone(ctx, existingTx.Client(), id, tombstoneKey) } tx, err := r.client.Tx(ctx) @@ -348,7 +346,7 @@ func (r *apiKeyRepository) DeleteWithAudit(ctx context.Context, id int64) error exec = tx.Client() } - if err := r.deleteWithAudit(ctx, exec, id, tombstoneKey); err != nil { + if err := r.deleteWithTombstone(ctx, exec, id, tombstoneKey); err != nil { return err } @@ -358,17 +356,7 @@ func (r *apiKeyRepository) DeleteWithAudit(ctx context.Context, id int64) error return nil } -func (r *apiKeyRepository) deleteWithAudit(ctx context.Context, exec *dbent.Client, id int64, tombstoneKey string) error { - // 1. 审计:数据源即 api_keys 当前行;WHERE deleted_at IS NULL 保证只对未删除行写一次。 - if _, err := exec.ExecContext(ctx, ` - INSERT INTO deleted_api_key_audits (key, api_key_id, user_id, key_name, deleted_at) - SELECT key, id, user_id, name, NOW() - FROM api_keys - WHERE id = $1 AND deleted_at IS NULL`, id); err != nil { - return err - } - - // 2. 软删除(tombstone 覆盖 key)。 +func (r *apiKeyRepository) deleteWithTombstone(ctx context.Context, exec *dbent.Client, id int64, tombstoneKey string) error { res, err := exec.ExecContext(ctx, ` UPDATE api_keys SET key = $1, deleted_at = NOW(), updated_at = NOW() diff --git a/backend/internal/repository/api_key_repo_integration_test.go b/backend/internal/repository/api_key_repo_integration_test.go index fdf9bc832..64a5e9933 100644 --- a/backend/internal/repository/api_key_repo_integration_test.go +++ b/backend/internal/repository/api_key_repo_integration_test.go @@ -556,7 +556,7 @@ func TestIncrementQuotaUsed_Concurrent(t *testing.T) { "并发递增后总和应为 %v,实际为 %v", float64(goroutines)*increment, got.QuotaUsed) } -func (s *APIKeyRepoSuite) TestDeleteWithAudit_WritesAuditAndSoftDeletes() { +func (s *APIKeyRepoSuite) TestDeleteWithAudit_TombstonesWithoutRetainingCredential() { user := s.mustCreateUser("delwithaudit@test.com") key := &service.APIKey{ UserID: user.ID, @@ -571,18 +571,24 @@ func (s *APIKeyRepoSuite) TestDeleteWithAudit_WritesAuditAndSoftDeletes() { _, err := s.repo.GetByID(s.ctx, key.ID) s.Require().Error(err) - rows, qErr := s.client.QueryContext(s.ctx, - `SELECT key, key_name, user_id, api_key_id FROM deleted_api_key_audits WHERE api_key_id = $1`, key.ID) - s.Require().NoError(qErr) - defer rows.Close() - s.Require().True(rows.Next(), "expected one audit row") - var auditKey, auditName string - var auditUserID, auditAPIKeyID int64 - s.Require().NoError(rows.Scan(&auditKey, &auditName, &auditUserID, &auditAPIKeyID)) - s.Require().Equal("sk-del-audit-1", auditKey) - s.Require().Equal("Audit Me", auditName) - s.Require().Equal(user.ID, auditUserID) - s.Require().Equal(key.ID, auditAPIKeyID) + var tombstone string + var deletedAt time.Time + rows, err := s.repo.sql.QueryContext(s.ctx, `SELECT key, deleted_at FROM api_keys WHERE id = $1`, key.ID) + s.Require().NoError(err) + s.Require().True(rows.Next()) + s.Require().NoError(rows.Scan(&tombstone, &deletedAt)) + s.Require().NoError(rows.Close()) + s.Require().NotEqual("sk-del-audit-1", tombstone) + s.Require().Contains(tombstone, "__deleted__") + + var auditCount int + auditRows, err := s.repo.sql.QueryContext(s.ctx, + `SELECT COUNT(*) FROM deleted_api_key_audits WHERE api_key_id = $1`, key.ID) + s.Require().NoError(err) + s.Require().True(auditRows.Next()) + s.Require().NoError(auditRows.Scan(&auditCount)) + s.Require().NoError(auditRows.Close()) + s.Require().Zero(auditCount, "deleted credentials must not be retained") } func (s *APIKeyRepoSuite) TestDeleteWithAudit_RepeatIsIdempotent() { diff --git a/backend/internal/repository/auth_cache_invalidation_outbox_integration_test.go b/backend/internal/repository/auth_cache_invalidation_outbox_integration_test.go new file mode 100644 index 000000000..5e8cc286f --- /dev/null +++ b/backend/internal/repository/auth_cache_invalidation_outbox_integration_test.go @@ -0,0 +1,120 @@ +//go:build integration + +package repository + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "fmt" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/stretchr/testify/require" +) + +func TestAuthCacheInvalidationTriggers_CoverSecurityMutationsOnly(t *testing.T) { + ctx := context.Background() + suffix := time.Now().UnixNano() + group := mustCreateGroup(t, integrationEntClient, &service.Group{ + Name: fmt.Sprintf("auth-outbox-group-%d", suffix), RateMultiplier: 1, IsExclusive: true, + }) + user := mustCreateUser(t, integrationEntClient, &service.User{ + Email: fmt.Sprintf("auth-outbox-%d@example.com", suffix), Concurrency: 5, + }) + groupID := group.ID + keyValue := fmt.Sprintf("sk-auth-outbox-%d", suffix) + apiKeyRepo := NewAPIKeyRepository(integrationEntClient, integrationDB) + key := &service.APIKey{UserID: user.ID, GroupID: &groupID, Key: keyValue, Name: "outbox", Status: service.StatusActive} + require.NoError(t, apiKeyRepo.Create(ctx, key)) + + sum := sha256.Sum256([]byte(keyValue)) + cacheKey := hex.EncodeToString(sum[:]) + clear := func() { + _, err := integrationDB.ExecContext(ctx, "DELETE FROM auth_cache_invalidation_outbox WHERE cache_key = $1", cacheKey) + require.NoError(t, err) + } + count := func() int { + var value int + require.NoError(t, integrationDB.QueryRowContext(ctx, + "SELECT COUNT(*) FROM auth_cache_invalidation_outbox WHERE cache_key = $1", cacheKey).Scan(&value)) + return value + } + clear() + t.Cleanup(clear) + t.Cleanup(func() { + // Keep the shared integration database isolated for suites that assert + // platform-wide group counts. The final clear cleanup runs after this one + // and removes invalidations emitted by these hard deletes. + _, err := integrationDB.ExecContext(ctx, "DELETE FROM user_allowed_groups WHERE user_id = $1 OR group_id = $2", user.ID, group.ID) + require.NoError(t, err) + _, err = integrationDB.ExecContext(ctx, "DELETE FROM api_keys WHERE id = $1", key.ID) + require.NoError(t, err) + _, err = integrationDB.ExecContext(ctx, "DELETE FROM users WHERE id = $1", user.ID) + require.NoError(t, err) + _, err = integrationDB.ExecContext(ctx, "DELETE FROM groups WHERE id = $1", group.ID) + require.NoError(t, err) + }) + + _, err := integrationDB.ExecContext(ctx, ` + UPDATE api_keys + SET quota_used = quota_used + 1, + usage_5h = usage_5h + 1, + last_used_at = NOW() + WHERE id = $1`, key.ID) + require.NoError(t, err) + require.Zero(t, count(), "usage-only key updates must not enqueue") + + _, err = integrationDB.ExecContext(ctx, "UPDATE api_keys SET status = 'disabled' WHERE id = $1", key.ID) + require.NoError(t, err) + require.Equal(t, 1, count(), "key disable must enqueue") + clear() + _, err = integrationDB.ExecContext(ctx, "UPDATE api_keys SET status = 'active' WHERE id = $1", key.ID) + require.NoError(t, err) + clear() + + userRepo := NewUserRepository(integrationEntClient, integrationDB) + loadedUser, err := userRepo.GetByID(ctx, user.ID) + require.NoError(t, err) + loadedUser.Balance += 10 + require.NoError(t, userRepo.Update(ctx, loadedUser)) + require.Zero(t, count(), "balance update with unchanged allowed groups must not enqueue") + + _, err = integrationDB.ExecContext(ctx, "UPDATE users SET status = 'disabled' WHERE id = $1", user.ID) + require.NoError(t, err) + require.Equal(t, 1, count(), "user disable must enqueue all active keys") + clear() + _, err = integrationDB.ExecContext(ctx, "UPDATE users SET status = 'active' WHERE id = $1", user.ID) + require.NoError(t, err) + clear() + + _, err = integrationDB.ExecContext(ctx, "UPDATE groups SET name = name || '-cosmetic' WHERE id = $1", group.ID) + require.NoError(t, err) + require.Zero(t, count(), "cosmetic group update must not enqueue") + _, err = integrationDB.ExecContext(ctx, "UPDATE groups SET status = 'disabled' WHERE id = $1", group.ID) + require.NoError(t, err) + require.Equal(t, 1, count(), "group disable must enqueue bound keys") + clear() + _, err = integrationDB.ExecContext(ctx, "UPDATE groups SET status = 'active' WHERE id = $1", group.ID) + require.NoError(t, err) + clear() + + _, err = integrationDB.ExecContext(ctx, + "INSERT INTO user_allowed_groups (user_id, group_id) VALUES ($1, $2)", user.ID, group.ID) + require.NoError(t, err) + clear() + _, err = integrationDB.ExecContext(ctx, + "DELETE FROM user_allowed_groups WHERE user_id = $1 AND group_id = $2", user.ID, group.ID) + require.NoError(t, err) + require.Equal(t, 1, count(), "exclusive-group revocation must enqueue") + clear() + + require.NoError(t, apiKeyRepo.DeleteWithAudit(ctx, key.ID)) + require.Equal(t, 1, count(), "tombstone delete must hash OLD.key exactly once") + var stored string + require.NoError(t, integrationDB.QueryRowContext(ctx, + "SELECT cache_key FROM auth_cache_invalidation_outbox WHERE cache_key = $1 LIMIT 1", cacheKey).Scan(&stored)) + require.Equal(t, cacheKey, stored) + require.NotContains(t, stored, keyValue) +} diff --git a/backend/internal/repository/auth_cache_invalidation_outbox_repo.go b/backend/internal/repository/auth_cache_invalidation_outbox_repo.go new file mode 100644 index 000000000..fc77fd812 --- /dev/null +++ b/backend/internal/repository/auth_cache_invalidation_outbox_repo.go @@ -0,0 +1,159 @@ +package repository + +import ( + "context" + "database/sql" + "errors" + "fmt" + "strings" + "time" + + "github.com/Wei-Shaw/sub2api/internal/service" +) + +type authCacheInvalidationOutboxRepository struct { + db *sql.DB +} + +func NewAuthCacheInvalidationOutboxRepository(db *sql.DB) service.AuthCacheInvalidationOutboxRepository { + return &authCacheInvalidationOutboxRepository{db: db} +} + +func (r *authCacheInvalidationOutboxRepository) Claim(ctx context.Context, workerID string, limit int, lease time.Duration) ([]service.AuthCacheInvalidationEvent, error) { + if r == nil || r.db == nil { + return nil, errors.New("nil auth cache invalidation outbox database") + } + if limit <= 0 { + limit = 100 + } + leaseSeconds := int64(lease / time.Second) + if leaseSeconds < 1 { + leaseSeconds = 30 + } + rows, err := r.db.QueryContext(ctx, ` + WITH candidates AS ( + SELECT id + FROM auth_cache_invalidation_outbox + WHERE available_at <= NOW() + AND (claimed_at IS NULL OR claimed_at < NOW() - ($3 * INTERVAL '1 second')) + ORDER BY id ASC + LIMIT $2 + FOR UPDATE SKIP LOCKED + ) + UPDATE auth_cache_invalidation_outbox AS o + SET claimed_at = NOW(), claimed_by = $1 + FROM candidates AS c + WHERE o.id = c.id + RETURNING o.id, o.cache_key, o.attempts, o.delivery_stage, o.created_at + `, workerID, limit, leaseSeconds) + if err != nil { + return nil, err + } + defer func() { _ = rows.Close() }() + + events := make([]service.AuthCacheInvalidationEvent, 0, limit) + for rows.Next() { + var event service.AuthCacheInvalidationEvent + if err := rows.Scan(&event.ID, &event.CacheKey, &event.Attempts, &event.Stage, &event.CreatedAt); err != nil { + return nil, err + } + event.CacheKey = strings.TrimSpace(event.CacheKey) + events = append(events, event) + } + if err := rows.Err(); err != nil { + return nil, err + } + return events, nil +} + +func (r *authCacheInvalidationOutboxRepository) ScheduleSecondPass(ctx context.Context, id int64, workerID string, availableAt time.Time) error { + result, err := r.db.ExecContext(ctx, ` + UPDATE auth_cache_invalidation_outbox + SET delivery_stage = 1, + available_at = $3, + last_error = NULL, + claimed_at = NULL, + claimed_by = NULL + WHERE id = $1 AND claimed_by = $2 AND delivery_stage = 0 + `, id, workerID, availableAt) + if err != nil { + return err + } + affected, err := result.RowsAffected() + if err != nil { + return err + } + if affected != 1 { + return fmt.Errorf("auth cache invalidation claim %d cannot schedule second pass", id) + } + return nil +} + +func (r *authCacheInvalidationOutboxRepository) DeleteClaimed(ctx context.Context, id int64, workerID string) error { + result, err := r.db.ExecContext(ctx, ` + DELETE FROM auth_cache_invalidation_outbox + WHERE id = $1 AND claimed_by = $2 + `, id, workerID) + if err != nil { + return err + } + affected, err := result.RowsAffected() + if err != nil { + return err + } + if affected != 1 { + return fmt.Errorf("auth cache invalidation claim %d is no longer owned by %s", id, workerID) + } + return nil +} + +func (r *authCacheInvalidationOutboxRepository) RetryClaimed(ctx context.Context, id int64, workerID string, availableAt time.Time, lastError string) error { + result, err := r.db.ExecContext(ctx, ` + UPDATE auth_cache_invalidation_outbox + SET attempts = attempts + 1, + available_at = $3, + last_error = $4, + claimed_at = NULL, + claimed_by = NULL + WHERE id = $1 AND claimed_by = $2 + `, id, workerID, availableAt, lastError) + if err != nil { + return err + } + affected, err := result.RowsAffected() + if err != nil { + return err + } + if affected != 1 { + return fmt.Errorf("auth cache invalidation claim %d is no longer owned by %s", id, workerID) + } + return nil +} + +func (r *authCacheInvalidationOutboxRepository) Stats(ctx context.Context) (service.AuthCacheInvalidationOutboxStats, error) { + var ( + stats service.AuthCacheInvalidationOutboxStats + oldest sql.NullTime + lastError sql.NullString + ) + err := r.db.QueryRowContext(ctx, ` + SELECT COUNT(*), MIN(created_at), COALESCE(MAX(attempts), 0), + (SELECT last_error + FROM auth_cache_invalidation_outbox + WHERE last_error IS NOT NULL + ORDER BY available_at DESC, id DESC + LIMIT 1) + FROM auth_cache_invalidation_outbox + `).Scan(&stats.Pending, &oldest, &stats.MaxAttempts, &lastError) + if err != nil { + return stats, err + } + if oldest.Valid { + value := oldest.Time + stats.OldestCreatedAt = &value + } + if lastError.Valid { + stats.LastError = lastError.String + } + return stats, nil +} diff --git a/backend/internal/repository/auth_cache_invalidation_outbox_repo_test.go b/backend/internal/repository/auth_cache_invalidation_outbox_repo_test.go new file mode 100644 index 000000000..1723055df --- /dev/null +++ b/backend/internal/repository/auth_cache_invalidation_outbox_repo_test.go @@ -0,0 +1,123 @@ +package repository + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "strings" + "testing" + "time" + + sqlmock "github.com/DATA-DOG/go-sqlmock" + "github.com/Wei-Shaw/sub2api/migrations" + "github.com/stretchr/testify/require" +) + +func TestAuthCacheInvalidationOutboxRepository_ClaimUsesLeaseAndSkipLocked(t *testing.T) { + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { _ = db.Close() }() + + created := time.Now().UTC() + mock.ExpectQuery("(?s)claimed_at < NOW\\(\\) - .*FOR UPDATE SKIP LOCKED.*RETURNING"). + WithArgs("worker-a", 100, int64(30)). + WillReturnRows(sqlmock.NewRows([]string{"id", "cache_key", "attempts", "delivery_stage", "created_at"}). + AddRow(int64(4), strings.Repeat("a", 64), 2, 1, created)) + + repo := NewAuthCacheInvalidationOutboxRepository(db) + events, err := repo.Claim(context.Background(), "worker-a", 100, 30*time.Second) + require.NoError(t, err) + require.Len(t, events, 1) + require.Equal(t, int64(4), events[0].ID) + require.Equal(t, 1, events[0].Stage) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestAuthCacheInvalidationOutboxRepository_ClaimIsBoundedByDefault(t *testing.T) { + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { _ = db.Close() }() + mock.ExpectQuery("(?s)FROM auth_cache_invalidation_outbox.*LIMIT \\$2.*SKIP LOCKED"). + WithArgs("worker", 100, int64(30)). + WillReturnRows(sqlmock.NewRows([]string{"id", "cache_key", "attempts", "delivery_stage", "created_at"})) + repo := NewAuthCacheInvalidationOutboxRepository(db) + _, err = repo.Claim(context.Background(), "worker", 0, 0) + require.NoError(t, err) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestAuthCacheInvalidationOutboxRepository_ClaimOwnershipTransitions(t *testing.T) { + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { _ = db.Close() }() + repo := NewAuthCacheInvalidationOutboxRepository(db) + + next := time.Now().UTC().Add(time.Minute) + mock.ExpectExec("UPDATE auth_cache_invalidation_outbox"). + WithArgs(int64(1), "worker", next). + WillReturnResult(sqlmock.NewResult(0, 1)) + require.NoError(t, repo.ScheduleSecondPass(context.Background(), 1, "worker", next)) + + retryAt := next.Add(time.Minute) + mock.ExpectExec("UPDATE auth_cache_invalidation_outbox"). + WithArgs(int64(2), "worker", retryAt, "publish failed"). + WillReturnResult(sqlmock.NewResult(0, 1)) + require.NoError(t, repo.RetryClaimed(context.Background(), 2, "worker", retryAt, "publish failed")) + + mock.ExpectExec("DELETE FROM auth_cache_invalidation_outbox"). + WithArgs(int64(3), "worker"). + WillReturnResult(sqlmock.NewResult(0, 1)) + require.NoError(t, repo.DeleteClaimed(context.Background(), 3, "worker")) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestAuthCacheInvalidationOutboxRepository_RejectsLostClaim(t *testing.T) { + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { _ = db.Close() }() + mock.ExpectExec("DELETE FROM auth_cache_invalidation_outbox"). + WithArgs(int64(3), "old-worker"). + WillReturnResult(sqlmock.NewResult(0, 0)) + repo := NewAuthCacheInvalidationOutboxRepository(db) + err = repo.DeleteClaimed(context.Background(), 3, "old-worker") + require.ErrorContains(t, err, "no longer owned") +} + +func TestAuthCacheInvalidationOutboxRepository_StatsExposeDurableLagAndFailures(t *testing.T) { + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { _ = db.Close() }() + oldest := time.Now().UTC().Add(-time.Minute) + mock.ExpectQuery("(?s)SELECT COUNT\\(\\*\\), MIN\\(created_at\\), COALESCE\\(MAX\\(attempts\\), 0\\)"). + WillReturnRows(sqlmock.NewRows([]string{"count", "min", "max", "last_error"}).AddRow(5, oldest, 7, "redis down")) + repo := NewAuthCacheInvalidationOutboxRepository(db) + stats, err := repo.Stats(context.Background()) + require.NoError(t, err) + require.Equal(t, int64(5), stats.Pending) + require.Equal(t, 7, stats.MaxAttempts) + require.Equal(t, "redis down", stats.LastError) + require.NotNil(t, stats.OldestCreatedAt) +} + +func TestAuthCacheInvalidationMigration_SecurityCoverageAndNoPlaintextPayload(t *testing.T) { + content, err := migrations.FS.ReadFile("184_auth_cache_invalidation_outbox.sql") + require.NoError(t, err) + sqlText := string(content) + for _, required := range []string{ + "encode(sha256(convert_to(raw_key, 'UTF8')), 'hex')", + "OLD.key", "OLD.status", "OLD.deleted_at", "OLD.user_id", "OLD.group_id", + "OLD.ip_whitelist", "OLD.ip_blacklist", "OLD.expires_at", + "trg_users_auth_cache_invalidation", "trg_groups_auth_cache_invalidation", + "trg_user_allowed_groups_auth_cache_invalidation", "FOR EACH ROW", + "delivery_stage", "claimed_at", "available_at", + } { + require.Contains(t, sqlText, required) + } + require.NotContains(t, sqlText, "quota_used IS DISTINCT") + require.NotContains(t, sqlText, "last_used_at IS DISTINCT") + + plaintext := "sk-plaintext-must-not-be-stored" + sum := sha256.Sum256([]byte(plaintext)) + require.Len(t, hex.EncodeToString(sum[:]), 64) + require.NotContains(t, sqlText, plaintext) +} diff --git a/backend/internal/repository/http_upstream.go b/backend/internal/repository/http_upstream.go index 8e3cf60bb..ec42e189f 100644 --- a/backend/internal/repository/http_upstream.go +++ b/backend/internal/repository/http_upstream.go @@ -7,6 +7,7 @@ import ( "compress/gzip" "context" "crypto/tls" + "encoding/json" "errors" "fmt" "io" @@ -323,7 +324,7 @@ func (t *grokAccessDeniedFallbackTransport) RoundTrip(req *http.Request) (*http. } body, ok := bufferSmallResponseBody(resp, grokFallbackBodyLimit) - if !ok || !bytes.Contains(bytes.ToLower(body), []byte("access denied")) { + if !ok || !isGrokCLICompatibilityAccessDenied(body) { return resp, nil } @@ -350,6 +351,22 @@ func (t *grokAccessDeniedFallbackTransport) RoundTrip(req *http.Request) (*http. return fallbackResp, nil } +func isGrokCLICompatibilityAccessDenied(body []byte) bool { + lower := bytes.ToLower(body) + if bytes.Contains(lower, []byte("access denied")) { + return true + } + var payload struct { + Code string `json:"code"` + Error string `json:"error"` + } + if err := json.Unmarshal(body, &payload); err != nil || !strings.EqualFold(strings.TrimSpace(payload.Code), "permission_denied") { + return false + } + const chatEndpointDeniedPrefix = "access to the chat endpoint is denied. please ensure you're using the correct credentials. if you believe this is a mistake, please" + return strings.HasPrefix(strings.ToLower(strings.TrimSpace(payload.Error)), chatEndpointDeniedPrefix) +} + func isGrokCLIAccessDeniedFallbackCandidate(req *http.Request, resp *http.Response) bool { return req != nil && req.URL != nil && req.GetBody != nil && resp != nil && resp.StatusCode == http.StatusForbidden && diff --git a/backend/internal/repository/http_upstream_test.go b/backend/internal/repository/http_upstream_test.go index 77749b7a0..35df74799 100644 --- a/backend/internal/repository/http_upstream_test.go +++ b/backend/internal/repository/http_upstream_test.go @@ -311,6 +311,119 @@ func TestHTTPUpstreamDoFallsBackToOfficialGrokAPIOnCLIAccessDenied(t *testing.T) require.Empty(t, fallbackHeaders.Get("User-Agent")) } +func TestGrokAccessDeniedFallbackRecognizesChatEndpointPermissionDenied(t *testing.T) { + var hosts []string + transport := &grokAccessDeniedFallbackTransport{ + base: roundTripFunc(func(req *http.Request) (*http.Response, error) { + hosts = append(hosts, req.URL.Hostname()) + if req.URL.Hostname() == grokCLIProxyHost { + return &http.Response{ + StatusCode: http.StatusForbidden, + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader( + `{"code":"permission_denied","error":"Access to the chat endpoint is denied. Please ensure you're using the correct credentials. If you believe this is a mistake, please contact support."}`, + )), + Request: req, + }, nil + } + return &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader(`{"id":"response-ok"}`)), + Request: req, + }, nil + }), + } + + req, err := http.NewRequest(http.MethodPost, "https://cli-chat-proxy.grok.com/v1/responses", strings.NewReader(`{"model":"grok-4.5"}`)) + require.NoError(t, err) + req.Header.Set("Authorization", "Bearer oauth-token") + req.Header.Set("X-XAI-Token-Auth", "xai-grok-cli") + + resp, err := transport.RoundTrip(req) + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode) + require.NoError(t, resp.Body.Close()) + require.Equal(t, []string{grokCLIProxyHost, grokOfficialAPIHost}, hosts) +} + +func TestIsGrokCLICompatibilityAccessDenied(t *testing.T) { + tests := []struct { + name string + body string + want bool + }{ + {name: "legacy compatibility wording", body: `{"error":"Access denied"}`, want: true}, + { + name: "observed chat endpoint permission denial", + body: `{"code":"permission_denied","error":"Access to the chat endpoint is denied. Please ensure you're using the correct credentials. If you believe this is a mistake, please contact support."}`, + want: true, + }, + { + name: "entitlement denial using the same broad terms", + body: `{"code":"permission_denied","error":"Access to the chat endpoint is denied because a subscription is required"}`, + want: false, + }, + { + name: "different permission denied endpoint", + body: `{"code":"permission_denied","error":"Access to the billing endpoint is denied."}`, + want: false, + }, + { + name: "wrong structured error code", + body: `{"code":"subscription_required","error":"Access to the chat endpoint is denied. Please ensure you're using the correct credentials. If you believe this is a mistake, please contact support."}`, + want: false, + }, + {name: "malformed response", body: `permission_denied: chat endpoint denied`, want: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + require.Equal(t, tt.want, isGrokCLICompatibilityAccessDenied([]byte(tt.body))) + }) + } +} + +func TestIsGrokCLIAccessDeniedFallbackCandidateRequiresAuthenticatedReplayableCLI403(t *testing.T) { + newRequest := func() *http.Request { + req, err := http.NewRequest(http.MethodPost, "https://cli-chat-proxy.grok.com/v1/responses", strings.NewReader(`{"model":"grok-4.5"}`)) + require.NoError(t, err) + req.Header.Set("Authorization", "Bearer oauth-token") + req.Header.Set("X-XAI-Token-Auth", "xai-grok-cli") + return req + } + newResponse := func() *http.Response { return &http.Response{StatusCode: http.StatusForbidden} } + + t.Run("valid candidate", func(t *testing.T) { + require.True(t, isGrokCLIAccessDeniedFallbackCandidate(newRequest(), newResponse())) + }) + t.Run("non CLI host", func(t *testing.T) { + req := newRequest() + req.URL.Host = "api.x.ai" + require.False(t, isGrokCLIAccessDeniedFallbackCandidate(req, newResponse())) + }) + t.Run("missing CLI identity", func(t *testing.T) { + req := newRequest() + req.Header.Del("X-XAI-Token-Auth") + require.False(t, isGrokCLIAccessDeniedFallbackCandidate(req, newResponse())) + }) + t.Run("missing bearer authentication", func(t *testing.T) { + req := newRequest() + req.Header.Del("Authorization") + require.False(t, isGrokCLIAccessDeniedFallbackCandidate(req, newResponse())) + }) + t.Run("non forbidden response", func(t *testing.T) { + resp := newResponse() + resp.StatusCode = http.StatusUnauthorized + require.False(t, isGrokCLIAccessDeniedFallbackCandidate(newRequest(), resp)) + }) + t.Run("non replayable request", func(t *testing.T) { + req := newRequest() + req.GetBody = nil + require.False(t, isGrokCLIAccessDeniedFallbackCandidate(req, newResponse())) + }) +} + func TestHTTPUpstreamDoDoesNotFallbackForGrokEntitlementDenial(t *testing.T) { transport := &grokAccessDeniedFallbackTransport{ base: roundTripFunc(func(req *http.Request) (*http.Response, error) { diff --git a/backend/internal/repository/migrations_runner.go b/backend/internal/repository/migrations_runner.go index a071967f6..c5e13b25b 100644 --- a/backend/internal/repository/migrations_runner.go +++ b/backend/internal/repository/migrations_runner.go @@ -128,23 +128,31 @@ func applyMigrationsFS(ctx context.Context, db *sql.DB, fsys fs.FS) error { // 获取分布式锁,确保多实例部署时只有一个实例执行迁移。 // 这是 PostgreSQL 特有的 Advisory Lock 机制。 - if err := pgAdvisoryLock(ctx, db); err != nil { + lockConn, err := db.Conn(ctx) + if err != nil { + return fmt.Errorf("acquire migrations lock connection: %w", err) + } + defer func() { _ = lockConn.Close() }() + if err := pgAdvisoryLock(ctx, lockConn); err != nil { return err } defer func() { // 无论迁移是否成功,都要释放锁。 - // 使用 context.Background() 确保即使原 ctx 已取消也能释放锁。 - _ = pgAdvisoryUnlock(context.Background(), db) + // 独立超时确保原 ctx 取消后仍会尝试释放,但数据库链路异常不会 + // 无限阻塞进程退出。 + unlockCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _ = pgAdvisoryUnlock(unlockCtx, lockConn) }() // 创建迁移记录表(如果不存在)。 // 该表记录所有已应用的迁移及其校验和。 - if _, err := db.ExecContext(ctx, schemaMigrationsTableDDL); err != nil { + if _, err := lockConn.ExecContext(ctx, schemaMigrationsTableDDL); err != nil { return fmt.Errorf("create schema_migrations: %w", err) } // 自动对齐 Atlas 基线(如果检测到 legacy schema_migrations 且缺失 atlas_schema_revisions)。 - if err := ensureAtlasBaselineAligned(ctx, db, fsys); err != nil { + if err := ensureAtlasBaselineAligned(ctx, lockConn, fsys); err != nil { return err } @@ -175,7 +183,7 @@ func applyMigrationsFS(ctx context.Context, db *sql.DB, fsys fs.FS) error { // 检查该迁移是否已经应用 var existing string - rowErr := db.QueryRowContext(ctx, "SELECT checksum FROM schema_migrations WHERE filename = $1", name).Scan(&existing) + rowErr := lockConn.QueryRowContext(ctx, "SELECT checksum FROM schema_migrations WHERE filename = $1", name).Scan(&existing) if rowErr == nil { // 迁移已应用,验证校验和是否匹配 if existing != checksum { @@ -207,7 +215,7 @@ func applyMigrationsFS(ctx context.Context, db *sql.DB, fsys fs.FS) error { } if nonTx { - if err := prepareNonTransactionalMigration(ctx, db, name); err != nil { + if err := prepareNonTransactionalMigration(ctx, lockConn, name); err != nil { return fmt.Errorf("prepare migration %s: %w", name, err) } @@ -222,18 +230,18 @@ func applyMigrationsFS(ctx context.Context, db *sql.DB, fsys fs.FS) error { if stripSQLLineComment(trimmed) == "" { continue } - if _, err := db.ExecContext(ctx, trimmed); err != nil { + if _, err := lockConn.ExecContext(ctx, trimmed); err != nil { return fmt.Errorf("apply migration %s (non-tx statement %d): %w", name, i+1, err) } } - if _, err := db.ExecContext(ctx, "INSERT INTO schema_migrations (filename, checksum) VALUES ($1, $2)", name, checksum); err != nil { + if _, err := lockConn.ExecContext(ctx, "INSERT INTO schema_migrations (filename, checksum) VALUES ($1, $2)", name, checksum); err != nil { return fmt.Errorf("record migration %s (non-tx): %w", name, err) } continue } // 默认迁移在事务中执行,确保原子性:要么完全成功,要么完全回滚。 - tx, err := db.BeginTx(ctx, nil) + tx, err := lockConn.BeginTx(ctx, nil) if err != nil { return fmt.Errorf("begin migration %s: %w", name, err) } @@ -260,7 +268,14 @@ 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 { +type migrationConnection 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 + BeginTx(ctx context.Context, opts *sql.TxOptions) (*sql.Tx, error) +} + +func prepareNonTransactionalMigration(ctx context.Context, db migrationConnection, name string) error { switch name { case paymentOrdersOutTradeNoUniqueMigration: return preparePaymentOrdersOutTradeNoUniqueMigration(ctx, db) @@ -273,7 +288,7 @@ func prepareNonTransactionalMigration(ctx context.Context, db *sql.DB, name stri } } -func preparePaymentOrdersOutTradeNoUniqueMigration(ctx context.Context, db *sql.DB) error { +func preparePaymentOrdersOutTradeNoUniqueMigration(ctx context.Context, db migrationConnection) error { duplicates, err := findDuplicatePaymentOrderOutTradeNos(ctx, db) if err != nil { return fmt.Errorf("precheck duplicate out_trade_no: %w", err) @@ -289,7 +304,7 @@ func preparePaymentOrdersOutTradeNoUniqueMigration(ctx context.Context, db *sql. return dropInvalidIndexIfPresent(ctx, db, paymentOrdersOutTradeNoUniqueIndex) } -func dropInvalidIndexIfPresent(ctx context.Context, db *sql.DB, indexName string) error { +func dropInvalidIndexIfPresent(ctx context.Context, db migrationConnection, indexName string) error { invalid, err := indexIsInvalid(ctx, db, indexName) if err != nil { return fmt.Errorf("check invalid index %s: %w", indexName, err) @@ -304,7 +319,7 @@ func dropInvalidIndexIfPresent(ctx context.Context, db *sql.DB, indexName string return nil } -func findDuplicatePaymentOrderOutTradeNos(ctx context.Context, db *sql.DB) ([]string, error) { +func findDuplicatePaymentOrderOutTradeNos(ctx context.Context, db migrationConnection) ([]string, error) { rows, err := db.QueryContext(ctx, ` SELECT out_trade_no, COUNT(*) AS duplicate_count FROM payment_orders @@ -336,7 +351,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 migrationConnection, indexName string) (bool, error) { var invalid bool err := db.QueryRowContext(ctx, ` SELECT EXISTS ( @@ -352,7 +367,7 @@ func indexIsInvalid(ctx context.Context, db *sql.DB, indexName string) (bool, er return invalid, err } -func ensureAtlasBaselineAligned(ctx context.Context, db *sql.DB, fsys fs.FS) error { +func ensureAtlasBaselineAligned(ctx context.Context, db migrationConnection, fsys fs.FS) error { hasLegacy, err := tableExists(ctx, db, "schema_migrations") if err != nil { return fmt.Errorf("check schema_migrations: %w", err) @@ -393,7 +408,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 migrationConnection, tableName string) (bool, error) { var exists bool err := db.QueryRowContext(ctx, ` SELECT EXISTS ( @@ -524,7 +539,12 @@ func stripSQLLineComment(s string) string { // pgAdvisoryLock 获取 PostgreSQL Advisory Lock。 // Advisory Lock 是一种轻量级的锁机制,不与任何特定的数据库对象关联。 // 它非常适合用于应用层面的分布式锁场景,如迁移序列化。 -func pgAdvisoryLock(ctx context.Context, db *sql.DB) error { +type advisoryLockConnection interface { + QueryRowContext(ctx context.Context, query string, args ...any) *sql.Row + ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) +} + +func pgAdvisoryLock(ctx context.Context, db advisoryLockConnection) error { ticker := time.NewTicker(migrationsLockRetryInterval) defer ticker.Stop() @@ -546,7 +566,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 advisoryLockConnection) 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 6bb7914b9..4ab6274b9 100644 --- a/backend/internal/repository/migrations_runner_notx_test.go +++ b/backend/internal/repository/migrations_runner_notx_test.go @@ -275,6 +275,9 @@ func TestApplyMigrationsFS_TransactionalMigration(t *testing.T) { db, mock, err := sqlmock.New() require.NoError(t, err) defer func() { _ = db.Close() }() + // The advisory lock and all migration work must share one session. This also + // proves startup cannot self-deadlock when deployments cap the pool at one. + db.SetMaxOpenConns(1) prepareMigrationsBootstrapExpectations(mock) mock.ExpectQuery("SELECT checksum FROM schema_migrations WHERE filename = \\$1"). diff --git a/backend/internal/repository/migrations_schema_integration_test.go b/backend/internal/repository/migrations_schema_integration_test.go index 7ef98dd70..68eb790ed 100644 --- a/backend/internal/repository/migrations_schema_integration_test.go +++ b/backend/internal/repository/migrations_schema_integration_test.go @@ -5,11 +5,32 @@ package repository import ( "context" "database/sql" + "sync" "testing" + "time" "github.com/stretchr/testify/require" ) +func TestMigrationsRunner_ConcurrentInstancesSerializeOnSessionLock(t *testing.T) { + const instances = 2 + errorsByInstance := make([]error, instances) + var wg sync.WaitGroup + for i := 0; i < instances; i++ { + wg.Add(1) + go func(index int) { + defer wg.Done() + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + errorsByInstance[index] = ApplyMigrations(ctx, integrationDB) + }(i) + } + wg.Wait() + for i, err := range errorsByInstance { + require.NoErrorf(t, err, "migration instance %d", i) + } +} + func TestMigrationsRunner_IsIdempotent_AndSchemaIsUpToDate(t *testing.T) { tx := testTx(t) @@ -111,6 +132,13 @@ func TestMigrationsRunner_IsIdempotent_AndSchemaIsUpToDate(t *testing.T) { requireColumn(t, tx, "ops_system_logs", "api_key_id", "bigint", 0, true) requireIndex(t, tx, "ops_system_logs", "idx_ops_system_logs_api_key_id_created_at") + // Bounded ingress rejection security aggregates. + requireColumn(t, tx, "ops_ingress_reject_aggregates", "bucket_start", "timestamp with time zone", 0, false) + requireColumn(t, tx, "ops_ingress_reject_aggregates", "client_ip", "inet", 0, false) + requireColumn(t, tx, "ops_ingress_reject_aggregates", "request_count", "bigint", 0, false) + requireIndex(t, tx, "ops_ingress_reject_aggregates", "idx_ops_ingress_reject_aggregates_bucket") + requireIndex(t, tx, "ops_ingress_reject_aggregates", "idx_ops_ingress_reject_aggregates_ip_bucket") + // user_allowed_groups table should exist var uagRegclass sql.NullString require.NoError(t, tx.QueryRowContext(context.Background(), "SELECT to_regclass('public.user_allowed_groups')").Scan(&uagRegclass)) diff --git a/backend/internal/repository/ops_error_where_test.go b/backend/internal/repository/ops_error_where_test.go index c52461e8f..af9af28de 100644 --- a/backend/internal/repository/ops_error_where_test.go +++ b/backend/internal/repository/ops_error_where_test.go @@ -137,26 +137,17 @@ func TestBuildOpsErrorLogsWhere_CyberPolicyStatusExemption(t *testing.T) { } } -func TestBuildOpsErrorLogsWhere_MatchDeletedKeyOwner(t *testing.T) { +func TestBuildOpsErrorLogsWhere_UserOwnershipIsDirectOnly(t *testing.T) { uid := int64(42) - - // 开关开启 → 归属放宽为 OR(user_id 或 deleted_key_owner_user_id),且共用同一占位符 - on := &service.OpsErrorLogFilter{UserID: &uid, MatchDeletedKeyOwner: true} - whereOn, argsOn := buildOpsErrorLogsWhere(on) - if !strings.Contains(whereOn, "(e.user_id = $1 OR e.deleted_key_owner_user_id = $1)") { - t.Fatalf("MatchDeletedKeyOwner=true should widen to OR, got: %s", whereOn) + filter := &service.OpsErrorLogFilter{UserID: &uid} + where, args := buildOpsErrorLogsWhere(filter) + if !strings.Contains(where, "e.user_id = $1") { + t.Fatalf("user scope should match user_id exactly, got: %s", where) } - if len(argsOn) != 1 || argsOn[0] != uid { - t.Fatalf("expected single reused arg %d, got %v", uid, argsOn) + if len(args) != 1 || args[0] != uid { + t.Fatalf("expected user id arg %d, got %v", uid, args) } - - // 开关关闭(默认)→ 仅精确 user_id,绝不出现 deleted_key_owner_user_id(admin 回归) - off := &service.OpsErrorLogFilter{UserID: &uid} - whereOff, _ := buildOpsErrorLogsWhere(off) - if !strings.Contains(whereOff, "e.user_id = $1") { - t.Fatalf("default should match user_id exactly, got: %s", whereOff) - } - if strings.Contains(whereOff, "deleted_key_owner_user_id") { - t.Fatalf("default must NOT include deleted_key_owner_user_id, got: %s", whereOff) + if strings.Contains(where, "deleted_key_owner_user_id") { + t.Fatalf("user ownership must not depend on deleted-key attribution: %s", where) } } diff --git a/backend/internal/repository/ops_ingress_reject_repo.go b/backend/internal/repository/ops_ingress_reject_repo.go new file mode 100644 index 000000000..14625781c --- /dev/null +++ b/backend/internal/repository/ops_ingress_reject_repo.go @@ -0,0 +1,158 @@ +package repository + +import ( + "context" + "fmt" + "strings" + + "github.com/Wei-Shaw/sub2api/internal/service" +) + +const ingressRejectUpsertChunkSize = 500 + +func (r *opsRepository) BatchUpsertIngressRejects(ctx context.Context, items []*service.OpsIngressRejectAggregate) error { + if r == nil || r.db == nil || len(items) == 0 { + return nil + } + tx, err := r.db.BeginTx(ctx, nil) + if err != nil { + return err + } + defer func() { _ = tx.Rollback() }() + + for start := 0; start < len(items); start += ingressRejectUpsertChunkSize { + end := start + ingressRejectUpsertChunkSize + if end > len(items) { + end = len(items) + } + valid := make([]*service.OpsIngressRejectAggregate, 0, end-start) + for _, item := range items[start:end] { + if item != nil && item.RequestCount > 0 { + valid = append(valid, item) + } + } + if len(valid) == 0 { + continue + } + + var query strings.Builder + _, _ = query.WriteString(`INSERT INTO ops_ingress_reject_aggregates + (bucket_start, reject_reason, route_family, protocol, client_ip, user_id, api_key_id, request_count, first_seen, last_seen) +VALUES `) + args := make([]any, 0, len(valid)*10) + for i, item := range valid { + if i > 0 { + _ = query.WriteByte(',') + } + base := len(args) + fmt.Fprintf(&query, "($%d,$%d,$%d,$%d,$%d,$%d,$%d,$%d,$%d,$%d)", + base+1, base+2, base+3, base+4, base+5, base+6, base+7, base+8, base+9, base+10) + var userID, apiKeyID int64 + if item.UserID != nil { + userID = *item.UserID + } + if item.APIKeyID != nil { + apiKeyID = *item.APIKeyID + } + args = append(args, item.BucketStart.UTC(), item.RejectReason, item.RouteFamily, item.Protocol, + item.ClientIP, userID, apiKeyID, item.RequestCount, item.FirstSeen.UTC(), item.LastSeen.UTC()) + } + _, _ = query.WriteString(` +ON CONFLICT (bucket_start, reject_reason, route_family, protocol, client_ip, user_id, api_key_id) +DO UPDATE SET request_count = ops_ingress_reject_aggregates.request_count + EXCLUDED.request_count, + first_seen = LEAST(ops_ingress_reject_aggregates.first_seen, EXCLUDED.first_seen), + last_seen = GREATEST(ops_ingress_reject_aggregates.last_seen, EXCLUDED.last_seen), + updated_at = NOW()`) + if _, err := tx.ExecContext(ctx, query.String(), args...); err != nil { + return err + } + } + return tx.Commit() +} + +func (r *opsRepository) ListIngressRejects(ctx context.Context, filter *service.OpsIngressRejectFilter) (*service.OpsIngressRejectList, error) { + if r == nil || r.db == nil { + return nil, fmt.Errorf("nil ops repository") + } + if filter == nil { + filter = &service.OpsIngressRejectFilter{} + } + page, pageSize := filter.Page, filter.PageSize + if page <= 0 { + page = 1 + } + if pageSize <= 0 { + pageSize = 50 + } + if pageSize > 200 { + pageSize = 200 + } + + clauses := []string{"1=1"} + args := make([]any, 0) + add := func(expr string, value any) { + args = append(args, value) + clauses = append(clauses, fmt.Sprintf(expr, len(args))) + } + if filter.StartTime != nil { + add("bucket_start >= $%d", filter.StartTime.UTC()) + } + if filter.EndTime != nil { + add("bucket_start < $%d", filter.EndTime.UTC()) + } + if value := strings.TrimSpace(filter.RejectReason); value != "" { + add("reject_reason = $%d", value) + } + if value := strings.TrimSpace(filter.RouteFamily); value != "" { + add("route_family = $%d", value) + } + if value := strings.TrimSpace(filter.Protocol); value != "" { + add("protocol = $%d", value) + } + if value := strings.TrimSpace(filter.ClientIP); value != "" { + add("client_ip = $%d::inet", value) + } + if filter.UserID != nil { + add("user_id = $%d", *filter.UserID) + } + if filter.APIKeyID != nil { + add("api_key_id = $%d", *filter.APIKeyID) + } + where := "WHERE " + strings.Join(clauses, " AND ") + + var total int + if err := r.db.QueryRowContext(ctx, "SELECT COUNT(*) FROM ops_ingress_reject_aggregates "+where, args...).Scan(&total); err != nil { + return nil, err + } + args = append(args, pageSize, (page-1)*pageSize) + query := fmt.Sprintf(`SELECT id,bucket_start,reject_reason,route_family,protocol,host(client_ip),user_id,api_key_id,request_count,first_seen,last_seen +FROM ops_ingress_reject_aggregates %s ORDER BY bucket_start DESC,id DESC LIMIT $%d OFFSET $%d`, where, len(args)-1, len(args)) + rows, err := r.db.QueryContext(ctx, query, args...) + if err != nil { + return nil, err + } + defer func() { _ = rows.Close() }() + + result := &service.OpsIngressRejectList{ + Items: make([]*service.OpsIngressRejectAggregate, 0, pageSize), Total: total, Page: page, PageSize: pageSize, + } + for rows.Next() { + item := &service.OpsIngressRejectAggregate{} + var userID, apiKeyID int64 + if err := rows.Scan(&item.ID, &item.BucketStart, &item.RejectReason, &item.RouteFamily, &item.Protocol, + &item.ClientIP, &userID, &apiKeyID, &item.RequestCount, &item.FirstSeen, &item.LastSeen); err != nil { + return nil, err + } + if userID > 0 { + item.UserID = &userID + } + if apiKeyID > 0 { + item.APIKeyID = &apiKeyID + } + result.Items = append(result.Items, item) + } + if err := rows.Err(); err != nil { + return nil, err + } + return result, nil +} diff --git a/backend/internal/repository/ops_ingress_reject_repo_test.go b/backend/internal/repository/ops_ingress_reject_repo_test.go new file mode 100644 index 000000000..fb437118d --- /dev/null +++ b/backend/internal/repository/ops_ingress_reject_repo_test.go @@ -0,0 +1,32 @@ +package repository + +import ( + "context" + "testing" + "time" + + "github.com/DATA-DOG/go-sqlmock" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/stretchr/testify/require" +) + +func TestBatchUpsertIngressRejectsUsesFixedMultiRowChunks(t *testing.T) { + db, mock, err := sqlmock.New() + require.NoError(t, err) + t.Cleanup(func() { _ = db.Close() }) + repo := &opsRepository{db: db} + now := time.Now().UTC().Truncate(time.Minute) + items := make([]*service.OpsIngressRejectAggregate, ingressRejectUpsertChunkSize+1) + for i := range items { + items[i] = &service.OpsIngressRejectAggregate{ + BucketStart: now, RejectReason: "invalid_api_key", RouteFamily: "messages", + Protocol: "anthropic", ClientIP: "192.0.2.1", RequestCount: 1, FirstSeen: now, LastSeen: now, + } + } + mock.ExpectBegin() + mock.ExpectExec("INSERT INTO ops_ingress_reject_aggregates").WillReturnResult(sqlmock.NewResult(0, int64(ingressRejectUpsertChunkSize))) + mock.ExpectExec("INSERT INTO ops_ingress_reject_aggregates").WillReturnResult(sqlmock.NewResult(0, 1)) + mock.ExpectCommit() + require.NoError(t, repo.BatchUpsertIngressRejects(context.Background(), items)) + require.NoError(t, mock.ExpectationsWereMet()) +} diff --git a/backend/internal/repository/ops_repo.go b/backend/internal/repository/ops_repo.go index 635bcd1a1..56fec9a93 100644 --- a/backend/internal/repository/ops_repo.go +++ b/backend/internal/repository/ops_repo.go @@ -4,7 +4,6 @@ import ( "context" "database/sql" "encoding/json" - "errors" "fmt" "strings" "time" @@ -56,12 +55,9 @@ INSERT INTO ops_error_logs ( response_latency_ms, time_to_first_token_ms, created_at, - attempted_key_prefix, - deleted_key_owner_user_id, - deleted_key_name, api_key_prefix ) VALUES ( - $1,$2,$3,$4,$5,$6,$7,$8,$9,$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 + $1,$2,$3,$4,$5,$6,$7,$8,$9,$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 )` func NewOpsRepository(db *sql.DB) service.OpsRepository { @@ -170,9 +166,6 @@ func opsInsertErrorLogArgs(input *service.OpsInsertErrorLogInput) []any { opsNullInt64(input.ResponseLatencyMs), opsNullInt64(input.TimeToFirstTokenMs), input.CreatedAt, - opsNullString(input.AttemptedKeyPrefix), - opsNullInt64(input.DeletedKeyOwnerUserID), - opsNullString(input.DeletedKeyName), opsNullString(input.APIKeyPrefix), } } @@ -274,16 +267,12 @@ SELECT COALESCE(e.user_agent, ''), e.request_type, COALESCE(ak.name, ''), - ak.deleted_at, - COALESCE(e.deleted_key_name, ''), - e.deleted_key_owner_user_id, - COALESCE(du.email, '') + ak.deleted_at FROM ops_error_logs e LEFT JOIN accounts a ON e.account_id = a.id LEFT JOIN groups g ON e.group_id = g.id LEFT JOIN users u ON e.user_id = u.id LEFT JOIN users u2 ON e.resolved_by_user_id = u2.id -LEFT JOIN users du ON e.deleted_key_owner_user_id = du.id LEFT JOIN api_keys ak ON ak.id = e.api_key_id ` + where + ` ORDER BY ` + opsErrorLogsOrderBy(filter) + ` @@ -313,9 +302,6 @@ LIMIT $` + itoa(len(args)+1) + ` OFFSET $` + itoa(len(args)+2) var requestType sql.NullInt64 var apiKeyName string var apiKeyDeletedAt sql.NullTime - var deletedKeyName string - var deletedKeyOwnerID sql.NullInt64 - var deletedKeyOwnerEmail string if err := rows.Scan( &item.ID, &item.CreatedAt, @@ -352,9 +338,6 @@ LIMIT $` + itoa(len(args)+1) + ` OFFSET $` + itoa(len(args)+2) &requestType, &apiKeyName, &apiKeyDeletedAt, - &deletedKeyName, - &deletedKeyOwnerID, - &deletedKeyOwnerEmail, ); err != nil { return nil, err } @@ -395,21 +378,8 @@ LIMIT $` + itoa(len(args)+1) + ` OFFSET $` + itoa(len(args)+2) v := int16(requestType.Int64) item.RequestType = &v } - // Key 名称:优先关联到的 ak.name(已软删的 key name 仍保留); - // 关联不到(api_key_id 为空 / 历史硬删)时回退错误记录里快照的 deleted_key_name。 - if apiKeyName != "" { - item.APIKeyName = apiKeyName - } else { - item.APIKeyName = deletedKeyName - } - // 已删除:ak.deleted_at 非空(软删),或仅命中 deleted_key_name 兜底。 - item.APIKeyDeleted = apiKeyDeletedAt.Valid || (apiKeyName == "" && deletedKeyName != "") - // 已删除 KEY 所有者快照:认证失败行 user_id 为空,列表用户列以此回退。 - if deletedKeyOwnerID.Valid { - v := deletedKeyOwnerID.Int64 - item.DeletedKeyOwnerUserID = &v - item.DeletedKeyOwnerEmail = deletedKeyOwnerEmail - } + item.APIKeyName = apiKeyName + item.APIKeyDeleted = apiKeyDeletedAt.Valid out = append(out, &item) } if err := rows.Err(); err != nil { @@ -477,10 +447,6 @@ SELECT e.upstream_latency_ms, e.response_latency_ms, e.time_to_first_token_ms, - COALESCE(e.attempted_key_prefix, ''), - e.deleted_key_owner_user_id, - COALESCE(du.email, ''), - COALESCE(e.deleted_key_name, ''), COALESCE(e.api_key_prefix, ''), COALESCE(ak.name, ''), ak.deleted_at @@ -488,7 +454,6 @@ FROM ops_error_logs e LEFT JOIN users u ON e.user_id = u.id LEFT JOIN accounts a ON e.account_id = a.id LEFT JOIN groups g ON e.group_id = g.id -LEFT JOIN users du ON e.deleted_key_owner_user_id = du.id LEFT JOIN api_keys ak ON ak.id = e.api_key_id WHERE e.id = $1 LIMIT 1` @@ -509,7 +474,6 @@ LIMIT 1` var responseLatency sql.NullInt64 var ttft sql.NullInt64 var requestType sql.NullInt64 - var deletedKeyOwnerUserID sql.NullInt64 var detailAPIKeyName string var detailAPIKeyDeletedAt sql.NullTime @@ -557,10 +521,6 @@ LIMIT 1` &upstreamLatency, &responseLatency, &ttft, - &out.AttemptedKeyPrefix, - &deletedKeyOwnerUserID, - &out.DeletedKeyOwnerEmail, - &out.DeletedKeyName, &out.APIKeyPrefix, &detailAPIKeyName, &detailAPIKeyDeletedAt, @@ -626,18 +586,8 @@ LIMIT 1` v := int16(requestType.Int64) out.RequestType = &v } - if deletedKeyOwnerUserID.Valid { - v := deletedKeyOwnerUserID.Int64 - out.DeletedKeyOwnerUserID = &v - } - // Key 名称:优先关联到的 ak.name;关联不到时回退快照的 deleted_key_name。 - if detailAPIKeyName != "" { - out.APIKeyName = detailAPIKeyName - } else { - out.APIKeyName = out.DeletedKeyName - } - // 已删除:ak.deleted_at 非空(软删),或仅命中 deleted_key_name 兜底。 - out.APIKeyDeleted = detailAPIKeyDeletedAt.Valid || (detailAPIKeyName == "" && out.DeletedKeyName != "") + out.APIKeyName = detailAPIKeyName + out.APIKeyDeleted = detailAPIKeyDeletedAt.Valid // Normalize upstream_errors to empty string when stored as JSON null. out.UpstreamErrors = strings.TrimSpace(out.UpstreamErrors) @@ -648,26 +598,6 @@ LIMIT 1` return &out, nil } -// LookupDeletedKeyAudit 按明文 key 反查最近一条已删除 key 审计。 -// 同一 key 可能有多条历史(反复创建/删除),取 deleted_at 最近一条(id 作同毫秒 tiebreaker)。 -// 未命中返回 (nil, nil)。 -func (r *opsRepository) LookupDeletedKeyAudit(ctx context.Context, key string) (*service.DeletedKeyAuditResult, error) { - var res service.DeletedKeyAuditResult - err := r.db.QueryRowContext(ctx, ` - SELECT user_id, key_name - FROM deleted_api_key_audits - WHERE key = $1 - ORDER BY deleted_at DESC, id DESC - LIMIT 1`, key).Scan(&res.UserID, &res.KeyName) - if err != nil { - if errors.Is(err, sql.ErrNoRows) { - return nil, nil - } - return nil, err - } - return &res, nil -} - func (r *opsRepository) UpdateErrorResolution(ctx context.Context, errorID int64, resolved bool, resolvedByUserID *int64, resolvedAt *time.Time) error { if r == nil || r.db == nil { return fmt.Errorf("nil ops repository") @@ -1082,12 +1012,7 @@ func buildOpsErrorLogsWhere(filter *service.OpsErrorLogFilter) (string, []any) { if filter.UserID != nil && *filter.UserID > 0 { args = append(args, *filter.UserID) n := itoa(len(args)) - if filter.MatchDeletedKeyOwner { - // 用户侧:把「删 key 后认证失败」(user_id=NULL,靠 deleted_key_owner 归因)的记录也纳入。 - clauses = append(clauses, "(e.user_id = $"+n+" OR e.deleted_key_owner_user_id = $"+n+")") - } else { - clauses = append(clauses, "e.user_id = $"+n) - } + clauses = append(clauses, "e.user_id = $"+n) } if filter.APIKeyID != nil && *filter.APIKeyID > 0 { args = append(args, *filter.APIKeyID) diff --git a/backend/internal/repository/ops_repo_args_test.go b/backend/internal/repository/ops_repo_args_test.go index 5a88758f0..9a5864b46 100644 --- a/backend/internal/repository/ops_repo_args_test.go +++ b/backend/internal/repository/ops_repo_args_test.go @@ -14,7 +14,7 @@ func TestOpsInsertErrorLogArgsPreservesExplicitZeroUpstreamStatus(t *testing.T) zero := 0 args := opsInsertErrorLogArgs(&service.OpsInsertErrorLogInput{UpstreamStatusCode: &zero}) - require.Len(t, args, 41) + require.Len(t, args, 38) encoded, ok := args[27].(sql.NullInt64) require.True(t, ok) require.True(t, encoded.Valid) diff --git a/backend/internal/repository/ops_repo_get_error_log_by_id_integration_test.go b/backend/internal/repository/ops_repo_get_error_log_by_id_integration_test.go index 1411002be..dbc98cf84 100644 --- a/backend/internal/repository/ops_repo_get_error_log_by_id_integration_test.go +++ b/backend/internal/repository/ops_repo_get_error_log_by_id_integration_test.go @@ -11,47 +11,13 @@ import ( "github.com/stretchr/testify/require" ) -// TestGetErrorLogByID_DeletedKeyOwner 验证: -// 1. 带 deleted_key_owner_user_id 的记录能正确 JOIN users 返回 DeletedKeyOwnerEmail -// 2. 新列全为 NULL 的普通记录 Scan 不报错,这些字段为空/nil -func TestGetErrorLogByID_DeletedKeyOwner(t *testing.T) { +func TestGetErrorLogByID_APIKeyPrefixAndUpstreamStatus(t *testing.T) { ctx := context.Background() _, _ = integrationDB.ExecContext(ctx, "TRUNCATE ops_error_logs RESTART IDENTITY CASCADE") - repo := NewOpsRepository(integrationDB).(*opsRepository) - // ── Case 1: 带 deleted_key_owner 信息的记录 ────────────────────────────── - owner := mustCreateUser(t, integrationEntClient, &service.User{ - Email: "deleted-key-owner-" + time.Now().Format("150405.000000000") + "@example.com", - }) - - var insertedID int64 - err := integrationDB.QueryRowContext(ctx, ` - INSERT INTO ops_error_logs ( - error_phase, error_type, severity, status_code, created_at, - attempted_key_prefix, deleted_key_owner_user_id, deleted_key_name - ) VALUES ( - 'auth', 'INVALID_API_KEY', 'error', 401, NOW(), - 'sk-test-abc', $1, 'my-deleted-key' - ) RETURNING id`, - owner.ID, - ).Scan(&insertedID) - require.NoError(t, err) - require.Positive(t, insertedID) - - detail, err := repo.GetErrorLogByID(ctx, insertedID) - require.NoError(t, err) - require.NotNil(t, detail) - - require.Equal(t, "sk-test-abc", detail.AttemptedKeyPrefix) - require.NotNil(t, detail.DeletedKeyOwnerUserID) - require.Equal(t, owner.ID, *detail.DeletedKeyOwnerUserID) - require.Equal(t, owner.Email, detail.DeletedKeyOwnerEmail) - require.Equal(t, "my-deleted-key", detail.DeletedKeyName) - - // ── Case 2: 新列全为 NULL 的普通错误记录 ────────────────────────────────── var plainID int64 - err = integrationDB.QueryRowContext(ctx, ` + err := integrationDB.QueryRowContext(ctx, ` INSERT INTO ops_error_logs ( error_phase, error_type, severity, status_code, created_at ) VALUES ( @@ -59,20 +25,11 @@ func TestGetErrorLogByID_DeletedKeyOwner(t *testing.T) { ) RETURNING id`, ).Scan(&plainID) require.NoError(t, err) - require.Positive(t, plainID) plain, err := repo.GetErrorLogByID(ctx, plainID) require.NoError(t, err) - require.NotNil(t, plain) + require.Empty(t, plain.APIKeyPrefix) - require.Empty(t, plain.AttemptedKeyPrefix, "no prefix for plain error") - require.Nil(t, plain.DeletedKeyOwnerUserID, "no owner for plain error") - require.Empty(t, plain.DeletedKeyOwnerEmail, "no owner email for plain error") - require.Empty(t, plain.DeletedKeyName, "no key name for plain error") - require.Empty(t, plain.APIKeyPrefix, "no api key prefix for plain error") - - // ── Case 3: 有效(未删除)key 报错,经 InsertErrorLog 快照 api_key_prefix ────── - // 走真实 InsertErrorLog 写入路径(覆盖新列 + $41 占位符),再 GetErrorLogByID 读回。 validID, err := repo.InsertErrorLog(ctx, &service.OpsInsertErrorLogInput{ ErrorPhase: "request", ErrorType: "api_error", @@ -82,17 +39,11 @@ func TestGetErrorLogByID_DeletedKeyOwner(t *testing.T) { APIKeyPrefix: "sk-valid", }) require.NoError(t, err) - require.Positive(t, validID) valid, err := repo.GetErrorLogByID(ctx, validID) require.NoError(t, err) - require.NotNil(t, valid) - require.Equal(t, "sk-valid", valid.APIKeyPrefix) - require.Empty(t, valid.AttemptedKeyPrefix, "attempted prefix and api key prefix are mutually exclusive") - require.Nil(t, valid.DeletedKeyOwnerUserID, "valid key error has no deleted owner") - // ── Case 4: account_auth with no inference attempt preserves explicit 0 ── zero := 0 credentialFailureID, err := repo.InsertErrorLog(ctx, &service.OpsInsertErrorLogInput{ ErrorPhase: "account_auth", diff --git a/backend/internal/repository/ops_repo_lookup_deleted_key_audit_integration_test.go b/backend/internal/repository/ops_repo_lookup_deleted_key_audit_integration_test.go deleted file mode 100644 index c77aefb9e..000000000 --- a/backend/internal/repository/ops_repo_lookup_deleted_key_audit_integration_test.go +++ /dev/null @@ -1,36 +0,0 @@ -//go:build integration - -package repository - -import ( - "context" - "testing" - "time" - - "github.com/stretchr/testify/require" -) - -func TestOpsRepositoryLookupDeletedKeyAudit(t *testing.T) { - ctx := context.Background() - _, _ = integrationDB.ExecContext(ctx, "TRUNCATE deleted_api_key_audits RESTART IDENTITY") - repo := NewOpsRepository(integrationDB).(*opsRepository) - - // 同一 key 两条审计,取最近一条(deleted_at DESC, id DESC) - _, err := integrationDB.ExecContext(ctx, ` - INSERT INTO deleted_api_key_audits (key, api_key_id, user_id, key_name, deleted_at) - VALUES ('sk-lookup-1', 10, 100, 'old', $1), - ('sk-lookup-1', 11, 200, 'new', $2)`, - time.Now().Add(-time.Hour), time.Now()) - require.NoError(t, err) - - res, err := repo.LookupDeletedKeyAudit(ctx, "sk-lookup-1") - require.NoError(t, err) - require.NotNil(t, res) - require.Equal(t, int64(200), res.UserID) - require.Equal(t, "new", res.KeyName) - - // 未命中返回 nil - miss, err := repo.LookupDeletedKeyAudit(ctx, "sk-never-existed") - require.NoError(t, err) - require.Nil(t, miss) -} diff --git a/backend/internal/repository/user_repo.go b/backend/internal/repository/user_repo.go index e739ab1f1..e08ce2f3a 100644 --- a/backend/internal/repository/user_repo.go +++ b/backend/internal/repository/user_repo.go @@ -1039,24 +1039,44 @@ func (r *userRepository) syncUserAllowedGroupsWithClient(ctx context.Context, cl return nil } - // Keep join table as the source of truth for reads. - if _, err := client.UserAllowedGroup.Delete().Where(userallowedgroup.UserIDEQ(userID)).Exec(ctx); err != nil { + existingRows, err := client.UserAllowedGroup.Query(). + Where(userallowedgroup.UserIDEQ(userID)). + All(ctx) + if err != nil { return err } - unique := make(map[int64]struct{}, len(groupIDs)) + desired := make(map[int64]struct{}, len(groupIDs)) for _, id := range groupIDs { if id <= 0 { continue } - unique[id] = struct{}{} + desired[id] = struct{}{} } - if len(unique) > 0 { - creates := make([]*dbent.UserAllowedGroupCreate, 0, len(unique)) - for groupID := range unique { + existing := make(map[int64]struct{}, len(existingRows)) + removed := make([]int64, 0) + for _, row := range existingRows { + existing[row.GroupID] = struct{}{} + if _, keep := desired[row.GroupID]; !keep { + removed = append(removed, row.GroupID) + } + } + if len(removed) > 0 { + if _, err := client.UserAllowedGroup.Delete(). + Where(userallowedgroup.UserIDEQ(userID), userallowedgroup.GroupIDIn(removed...)). + Exec(ctx); err != nil { + return err + } + } + + creates := make([]*dbent.UserAllowedGroupCreate, 0, len(desired)) + for groupID := range desired { + if _, present := existing[groupID]; !present { creates = append(creates, client.UserAllowedGroup.Create().SetUserID(userID).SetGroupID(groupID)) } + } + if len(creates) > 0 { if err := client.UserAllowedGroup. CreateBulk(creates...). OnConflictColumns(userallowedgroup.FieldUserID, userallowedgroup.FieldGroupID). diff --git a/backend/internal/repository/user_repo_delete_atomicity_integration_test.go b/backend/internal/repository/user_repo_delete_atomicity_integration_test.go index 5c1aecbe5..9a3ac580e 100644 --- a/backend/internal/repository/user_repo_delete_atomicity_integration_test.go +++ b/backend/internal/repository/user_repo_delete_atomicity_integration_test.go @@ -14,7 +14,7 @@ import ( ) // TestUserRepository_DeleteUser_AtomicWithAPIKeys 复现 AdminService.DeleteUser 的事务编排场景: -// 把"删 API Key"(apiKeyRepo.DeleteWithAudit) 与"删 User"(userRepo.Delete) 放进同一个外部事务时, +// 把"tombstone 并删 API Key"(apiKeyRepo.DeleteWithAudit) 与"删 User"(userRepo.Delete) 放进同一个外部事务时, // userRepo.Delete 必须复用 context 中的事务,而不是用 base client 自起一个独立事务并提前提交。 // // 用例用"回滚外层事务"来模拟 commit 失败 / 中止: @@ -91,5 +91,5 @@ func TestUserRepository_DeleteUser_AtomicWithAPIKeys(t *testing.T) { require.NoError(t, integrationDB.QueryRowContext(ctx, `SELECT COUNT(*) FROM deleted_api_key_audits WHERE user_id = $1`, user.ID).Scan(&auditCount)) - require.Equal(t, 2, auditCount, "提交后应为每个被删 Key 写入一行审计") + require.Zero(t, auditCount, "提交后也不得保留被删 Key 的凭据材料") } diff --git a/backend/internal/repository/wire.go b/backend/internal/repository/wire.go index 581831912..0bfd2d005 100644 --- a/backend/internal/repository/wire.go +++ b/backend/internal/repository/wire.go @@ -126,6 +126,7 @@ var ProviderSet = wire.NewSet( NewLeaderLockCache, ProvideSchedulerCache, NewSchedulerOutboxRepository, + NewAuthCacheInvalidationOutboxRepository, NewProxyLatencyCache, NewTotpCache, NewRefreshTokenCache, diff --git a/backend/internal/securityaudit/prompt_module.go b/backend/internal/securityaudit/prompt_module.go index 691cc7dbd..e1199fcfa 100644 --- a/backend/internal/securityaudit/prompt_module.go +++ b/backend/internal/securityaudit/prompt_module.go @@ -16,6 +16,7 @@ var ProviderSet = wire.NewSet( wire.Bind(new(ConfigStore), new(*ConfigManager)), NewPromptService, wire.Bind(new(PromptEngine), new(*PromptService)), + wire.Bind(new(PromptAdminService), new(*PromptService)), NewLegacyModerationAdapter, NewCoordinator, NewPromptAdminHandler, diff --git a/backend/internal/server/api_contract_test.go b/backend/internal/server/api_contract_test.go index 6b9165217..86ac46c2a 100644 --- a/backend/internal/server/api_contract_test.go +++ b/backend/internal/server/api_contract_test.go @@ -1876,6 +1876,10 @@ func (s *stubAccountRepo) ListSchedulableUngroupedByPlatforms(ctx context.Contex return nil, errors.New("not implemented") } +func (s *stubAccountRepo) ListModelAvailabilityCandidates(ctx context.Context, groupID *int64, platforms []string, includeGrouped bool) ([]service.Account, error) { + return nil, errors.New("not implemented") +} + func (s *stubAccountRepo) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error { return errors.New("not implemented") } diff --git a/backend/internal/server/http.go b/backend/internal/server/http.go index 89ab9bf48..65548c9e2 100644 --- a/backend/internal/server/http.go +++ b/backend/internal/server/http.go @@ -103,8 +103,9 @@ func ProvideRouter( func ProvideHTTPServer(cfg *config.Config, router *gin.Engine) *http.Server { httpHandler := http.Handler(router) server := &http.Server{ - Addr: cfg.Server.Address(), - Handler: httpHandler, + Addr: cfg.Server.Address(), + Handler: httpHandler, + MaxHeaderBytes: cfg.Server.MaxHeaderBytes, // ReadHeaderTimeout: 读取请求头的超时时间,防止慢速请求头攻击 ReadHeaderTimeout: time.Duration(cfg.Server.ReadHeaderTimeout) * time.Second, // IdleTimeout: 空闲连接超时时间,释放不活跃的连接资源 diff --git a/backend/internal/server/http_ingress_test.go b/backend/internal/server/http_ingress_test.go new file mode 100644 index 000000000..63cf7a047 --- /dev/null +++ b/backend/internal/server/http_ingress_test.go @@ -0,0 +1,126 @@ +//go:build unit + +package server + +import ( + "bufio" + "context" + "errors" + "io" + "net" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func ingressTestConfig() *config.Config { + return &config.Config{ + Server: config.ServerConfig{ + Host: "127.0.0.1", + ReadHeaderTimeout: 1, + IdleTimeout: 5, + MaxHeaderBytes: 8 * 1024, + MaxRequestBodySize: 1024, + }, + Gateway: config.GatewayConfig{MaxBodySize: 1024}, + } +} + +func TestProvideHTTPServerAppliesIngressLimits(t *testing.T) { + srv := ProvideHTTPServer(ingressTestConfig(), gin.New()) + require.Equal(t, 8*1024, srv.MaxHeaderBytes) + require.Equal(t, time.Second, srv.ReadHeaderTimeout) + require.Equal(t, 5*time.Second, srv.IdleTimeout) +} + +func TestProvideHTTPServerEnablesBoundedH2C(t *testing.T) { + cfg := ingressTestConfig() + cfg.Server.H2C = config.H2CConfig{ + Enabled: true, + MaxConcurrentStreams: 25, + IdleTimeout: 30, + MaxReadFrameSize: 64 * 1024, + MaxUploadBufferPerConnection: 1024 * 1024, + MaxUploadBufferPerStream: 256 * 1024, + } + srv := ProvideHTTPServer(cfg, gin.New()) + require.NotNil(t, srv.Protocols) + require.True(t, srv.Protocols.UnencryptedHTTP2()) + require.True(t, srv.Protocols.HTTP1()) +} + +func TestHTTPServerRejectsOversizedHTTP1Header(t *testing.T) { + r := gin.New() + r.GET("/", func(c *gin.Context) { c.Status(http.StatusOK) }) + srv := ProvideHTTPServer(ingressTestConfig(), r) + addr, stop := serveIngressTestServer(t, srv) + defer stop() + + conn, err := net.DialTimeout("tcp", addr, time.Second) + require.NoError(t, err) + defer func() { _ = conn.Close() }() + _ = conn.SetDeadline(time.Now().Add(3 * time.Second)) + _, err = io.WriteString(conn, "GET / HTTP/1.1\r\nHost: test\r\nX-Fill: "+strings.Repeat("a", 32*1024)+"\r\n\r\n") + require.NoError(t, err) + resp, err := http.ReadResponse(bufio.NewReader(conn), nil) + require.NoError(t, err) + defer func() { _ = resp.Body.Close() }() + require.Equal(t, http.StatusRequestHeaderFieldsTooLarge, resp.StatusCode) +} + +func TestHTTPServerClosesSlowIncompleteHeader(t *testing.T) { + r := gin.New() + r.GET("/", func(c *gin.Context) { c.Status(http.StatusOK) }) + srv := ProvideHTTPServer(ingressTestConfig(), r) + addr, stop := serveIngressTestServer(t, srv) + defer stop() + + conn, err := net.DialTimeout("tcp", addr, time.Second) + require.NoError(t, err) + defer func() { _ = conn.Close() }() + _, err = io.WriteString(conn, "GET / HTTP/1.1\r\nHost: test\r\nX-Slow:") + require.NoError(t, err) + time.Sleep(1200 * time.Millisecond) + _ = conn.SetReadDeadline(time.Now().Add(time.Second)) + _, err = bufio.NewReader(conn).ReadByte() + require.Error(t, err) +} + +func TestHTTPServerGlobalBodyLimit(t *testing.T) { + r := gin.New() + r.POST("/", func(c *gin.Context) { + _, err := io.ReadAll(c.Request.Body) + if err != nil { + var maxErr *http.MaxBytesError + if errors.As(err, &maxErr) { + c.Status(http.StatusRequestEntityTooLarge) + return + } + } + c.Status(http.StatusOK) + }) + srv := ProvideHTTPServer(ingressTestConfig(), r) + req, err := http.NewRequest(http.MethodPost, "/", strings.NewReader(strings.Repeat("x", 1025))) + require.NoError(t, err) + rec := httptest.NewRecorder() + srv.Handler.ServeHTTP(rec, req) + require.Equal(t, http.StatusRequestEntityTooLarge, rec.Code) +} + +func serveIngressTestServer(t *testing.T, srv *http.Server) (string, func()) { + t.Helper() + ln, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + go func() { _ = srv.Serve(ln) }() + return ln.Addr().String(), func() { + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + _ = srv.Shutdown(ctx) + } +} diff --git a/backend/internal/server/middleware/api_key_auth.go b/backend/internal/server/middleware/api_key_auth.go index b9013815a..bcec39760 100644 --- a/backend/internal/server/middleware/api_key_auth.go +++ b/backend/internal/server/middleware/api_key_auth.go @@ -15,6 +15,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)) @@ -32,10 +34,23 @@ 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 rejectInvalidAuthAbuse(c, apiKeyService) { + AbortWithError(c, http.StatusTooManyRequests, "INVALID_AUTH_RATE_LIMITED", "Too many invalid authentication attempts; retry later") + return + } + + if apiKeyHeadersTooLarge(c) { + recordInvalidAuthFailure(c, apiKeyService) + MarkIngressRejected(c, IngressRejectInvalidAPIKey) + AbortWithError(c, http.StatusUnauthorized, "INVALID_API_KEY", "Invalid API key") + return + } queryKey := strings.TrimSpace(c.Query("key")) queryApiKey := strings.TrimSpace(c.Query("api_key")) if queryKey != "" || queryApiKey != "" { + recordInvalidAuthFailure(c, apiKeyService) + MarkIngressRejected(c, IngressRejectQueryAPIKeyDeprecated) AbortWithError(c, 400, "api_key_in_query_deprecated", "API key in query parameter is deprecated. Please use Authorization header instead.") return } @@ -56,6 +71,12 @@ func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscripti if apiKeyString == "" { apiKeyString = c.GetHeader("x-api-key") } + if len(apiKeyString) > service.MaxAPIKeyCredentialBytes { + recordInvalidAuthFailure(c, apiKeyService) + MarkIngressRejected(c, IngressRejectInvalidAPIKey) + AbortWithError(c, http.StatusUnauthorized, "INVALID_API_KEY", "Invalid API key") + return + } // 如果x-api-key header中没有,尝试从x-goog-api-key header中提取(Gemini CLI兼容) if apiKeyString == "" { @@ -64,6 +85,12 @@ func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscripti // 如果所有header都没有API key if apiKeyString == "" { + recordInvalidAuthFailure(c, apiKeyService) + if hasAPIKeyCredentialInput(c) { + MarkIngressRejected(c, IngressRejectInvalidAPIKey) + } else { + MarkIngressRejected(c, IngressRejectAPIKeyRequired) + } AbortWithError(c, 401, "API_KEY_REQUIRED", "API key is required in Authorization header (Bearer scheme), x-api-key header, or x-goog-api-key header") return } @@ -73,9 +100,16 @@ func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscripti apiKey, err := apiKeyService.GetByKey(c.Request.Context(), apiKeyString) if err != nil { if errors.Is(err, service.ErrAPIKeyNotFound) { + recordInvalidAuthFailure(c, apiKeyService) + MarkIngressRejected(c, IngressRejectInvalidAPIKey) AbortWithError(c, 401, "INVALID_API_KEY", "Invalid API key") return } + if errors.Is(err, service.ErrAPIKeyAuthOverloaded) { + MarkIngressRejected(c, IngressRejectAPIKeyAuthOverloaded) + AbortWithError(c, http.StatusServiceUnavailable, "API_KEY_AUTH_OVERLOADED", "API key authentication is temporarily unavailable") + return + } AbortWithError(c, 500, "INTERNAL_ERROR", "Failed to validate API key") return } @@ -90,6 +124,7 @@ func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscripti if !apiKey.IsActive() && apiKey.Status != service.StatusAPIKeyExpired && apiKey.Status != service.StatusAPIKeyQuotaExhausted { + MarkIngressRejected(c, IngressRejectAPIKeyDisabled) AbortWithError(c, 401, "API_KEY_DISABLED", "API key is disabled") return } @@ -104,6 +139,7 @@ func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscripti clientIP = "unknown" } service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonIPRestriction) + MarkIngressRejected(c, IngressRejectIPRestricted) AbortWithError(c, 403, "ACCESS_DENIED", fmt.Sprintf("Access denied. Your IP is %s", clientIP)) return } @@ -117,6 +153,7 @@ func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscripti // 检查用户状态 if !apiKey.User.IsActive() { + MarkIngressRejected(c, IngressRejectUserInactive) AbortWithError(c, 401, "USER_INACTIVE", "User account is not active") return } @@ -250,6 +287,24 @@ func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscripti } } +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 +} + +func hasAPIKeyCredentialInput(c *gin.Context) bool { + if c == nil { + return false + } + return c.GetHeader("Authorization") != "" || + c.GetHeader("x-api-key") != "" || + c.GetHeader("x-goog-api-key") != "" +} + func isAsyncImageTaskRead(method, path string) bool { if method != http.MethodGet { return false @@ -321,6 +376,11 @@ func abortIfAPIKeyGroupUnavailable(c *gin.Context, apiKey *service.APIKey) bool return false } service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonAPIKeyGroupUnavailable) + if code == "GROUP_DELETED" { + MarkIngressRejected(c, IngressRejectGroupDeleted) + } else { + MarkIngressRejected(c, IngressRejectGroupDisabled) + } AbortWithError(c, 403, code, message) return true } @@ -330,6 +390,7 @@ func abortIfAPIKeyGroupNotAllowed(c *gin.Context, apiKey *service.APIKey) bool { return false } service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonAPIKeyGroupUnavailable) + MarkIngressRejected(c, IngressRejectGroupNotAllowed) AbortWithError(c, 403, "GROUP_NOT_ALLOWED", "API Key 所属专属分组不再允许当前用户使用") return true } diff --git a/backend/internal/server/middleware/api_key_auth_google.go b/backend/internal/server/middleware/api_key_auth_google.go index 49bf6b426..9efbbcb2c 100644 --- a/backend/internal/server/middleware/api_key_auth_google.go +++ b/backend/internal/server/middleware/api_key_auth_google.go @@ -24,22 +24,53 @@ func APIKeyAuthGoogle(apiKeyService *service.APIKeyService, cfg *config.Config) // It is intended for Gemini native endpoints (/v1beta) to match Gemini SDK expectations. func APIKeyAuthWithSubscriptionGoogle(apiKeyService *service.APIKeyService, subscriptionService *service.SubscriptionService, cfg *config.Config) gin.HandlerFunc { return func(c *gin.Context) { + if rejectInvalidAuthAbuse(c, apiKeyService) { + abortWithGoogleError(c, 429, "Too many invalid authentication attempts; retry later") + return + } + if apiKeyHeadersTooLarge(c) { + recordInvalidAuthFailure(c, apiKeyService) + MarkIngressRejected(c, IngressRejectInvalidAPIKey) + abortWithGoogleError(c, 401, "Invalid API key") + return + } if v := strings.TrimSpace(c.Query("api_key")); v != "" { + recordInvalidAuthFailure(c, apiKeyService) + MarkIngressRejected(c, IngressRejectQueryAPIKeyDeprecated) abortWithGoogleError(c, 400, "Query parameter api_key is deprecated. Use Authorization header or key instead.") return } apiKeyString := extractAPIKeyForGoogle(c) if apiKeyString == "" { + recordInvalidAuthFailure(c, apiKeyService) + if hasAPIKeyCredentialInput(c) { + MarkIngressRejected(c, IngressRejectInvalidAPIKey) + } else { + MarkIngressRejected(c, IngressRejectAPIKeyRequired) + } abortWithGoogleError(c, 401, "API key is required") return } + if len(apiKeyString) > service.MaxAPIKeyCredentialBytes { + recordInvalidAuthFailure(c, apiKeyService) + MarkIngressRejected(c, IngressRejectInvalidAPIKey) + abortWithGoogleError(c, 401, "Invalid API key") + return + } apiKey, err := apiKeyService.GetByKey(c.Request.Context(), apiKeyString) if err != nil { if errors.Is(err, service.ErrAPIKeyNotFound) { + recordInvalidAuthFailure(c, apiKeyService) + MarkIngressRejected(c, IngressRejectInvalidAPIKey) abortWithGoogleError(c, 401, "Invalid API key") return } + if errors.Is(err, service.ErrAPIKeyAuthOverloaded) { + MarkIngressRejected(c, IngressRejectAPIKeyAuthOverloaded) + abortWithGoogleError(c, 503, "API key authentication is temporarily unavailable") + return + } abortWithGoogleError(c, 500, "Failed to validate API key") return } @@ -53,6 +84,7 @@ func APIKeyAuthWithSubscriptionGoogle(apiKeyService *service.APIKeyService, subs if !apiKey.IsActive() && apiKey.Status != service.StatusAPIKeyExpired && apiKey.Status != service.StatusAPIKeyQuotaExhausted { + MarkIngressRejected(c, IngressRejectAPIKeyDisabled) abortWithGoogleError(c, 401, "API key is disabled") return } @@ -66,6 +98,7 @@ func APIKeyAuthWithSubscriptionGoogle(apiKeyService *service.APIKeyService, subs clientIP = "unknown" } service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonIPRestriction) + MarkIngressRejected(c, IngressRejectIPRestricted) abortWithGoogleError(c, 403, fmt.Sprintf("Access denied. Your IP is %s", clientIP)) return } @@ -76,17 +109,24 @@ func APIKeyAuthWithSubscriptionGoogle(apiKeyService *service.APIKeyService, subs return } if !apiKey.User.IsActive() { + MarkIngressRejected(c, IngressRejectUserInactive) abortWithGoogleError(c, 401, "User account is not active") return } - if _, message, ok := validateAPIKeyGroupAvailable(apiKey); !ok { + if code, message, ok := validateAPIKeyGroupAvailable(apiKey); !ok { service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonAPIKeyGroupUnavailable) + if code == "GROUP_DELETED" { + MarkIngressRejected(c, IngressRejectGroupDeleted) + } else { + MarkIngressRejected(c, IngressRejectGroupDisabled) + } abortWithGoogleError(c, 403, message) return } // 专属分组授权校验:用户对该专属分组的授权被撤销后应拒绝(与主中间件一致,防止越权)。 if !validateAPIKeyGroupAllowed(apiKey) { service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonAPIKeyGroupUnavailable) + MarkIngressRejected(c, IngressRejectGroupNotAllowed) abortWithGoogleError(c, 403, "API Key 所属专属分组不再允许当前用户使用") return } diff --git a/backend/internal/server/middleware/api_key_auth_google_test.go b/backend/internal/server/middleware/api_key_auth_google_test.go index 746d238ca..1b1e196a2 100644 --- a/backend/internal/server/middleware/api_key_auth_google_test.go +++ b/backend/internal/server/middleware/api_key_auth_google_test.go @@ -6,6 +6,8 @@ import ( "errors" "net/http" "net/http/httptest" + "strings" + "sync/atomic" "testing" "time" @@ -18,6 +20,59 @@ import ( "github.com/stretchr/testify/require" ) +func TestGoogleAPIKeyAuthRejectsOversizedCredentialsBeforeLookup(t *testing.T) { + gin.SetMode(gin.TestMode) + var calls atomic.Int32 + repo := fakeAPIKeyRepo{getByKey: func(context.Context, string) (*service.APIKey, error) { + calls.Add(1) + return nil, service.ErrAPIKeyNotFound + }} + cfg := &config.Config{RunMode: config.RunModeSimple} + svc := service.NewAPIKeyService(repo, nil, nil, nil, nil, nil, cfg) + r := gin.New() + var reason IngressRejectReason + var rejected bool + r.Use(func(c *gin.Context) { + c.Next() + reason, rejected = GetIngressRejectReason(c) + }) + r.Use(APIKeyAuthGoogle(svc, cfg)) + r.GET("/v1beta/models", func(c *gin.Context) { c.Status(http.StatusOK) }) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/v1beta/models", nil) + req.Header.Set("x-goog-api-key", strings.Repeat("x", service.MaxAPIKeyCredentialBytes+1)) + r.ServeHTTP(w, req) + require.Equal(t, http.StatusUnauthorized, w.Code) + require.Zero(t, calls.Load()) + require.True(t, rejected) + require.Equal(t, IngressRejectInvalidAPIKey, reason) +} + +func TestGoogleAPIKeyAuthMarksLookupBulkheadRejection(t *testing.T) { + gin.SetMode(gin.TestMode) + repo := fakeAPIKeyRepo{getByKey: func(context.Context, string) (*service.APIKey, error) { + return nil, service.ErrAPIKeyAuthOverloaded + }} + cfg := &config.Config{RunMode: config.RunModeSimple} + svc := service.NewAPIKeyService(repo, nil, nil, nil, nil, nil, cfg) + r := gin.New() + var reason IngressRejectReason + var rejected bool + r.Use(func(c *gin.Context) { + c.Next() + reason, rejected = GetIngressRejectReason(c) + }) + r.Use(APIKeyAuthGoogle(svc, cfg)) + r.GET("/v1beta/models", func(c *gin.Context) { c.Status(http.StatusOK) }) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/v1beta/models", nil) + req.Header.Set("x-goog-api-key", "valid-shape") + r.ServeHTTP(w, req) + require.Equal(t, http.StatusServiceUnavailable, w.Code) + require.True(t, rejected) + require.Equal(t, IngressRejectAPIKeyAuthOverloaded, reason) +} + type fakeAPIKeyRepo struct { getByKey func(ctx context.Context, key string) (*service.APIKey, error) updateLastUsed func(ctx context.Context, id int64, usedAt time.Time) error @@ -376,6 +431,12 @@ func TestApiKeyAuthWithSubscriptionGoogle_InvalidKey(t *testing.T) { return nil, service.ErrAPIKeyNotFound }, }) + var rejectReason IngressRejectReason + var rejected bool + r.Use(func(c *gin.Context) { + c.Next() + rejectReason, rejected = GetIngressRejectReason(c) + }) r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, &config.Config{})) r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) }) @@ -390,6 +451,8 @@ func TestApiKeyAuthWithSubscriptionGoogle_InvalidKey(t *testing.T) { require.Equal(t, http.StatusUnauthorized, resp.Error.Code) require.Equal(t, "Invalid API key", resp.Error.Message) require.Equal(t, "UNAUTHENTICATED", resp.Error.Status) + require.True(t, rejected) + require.Equal(t, IngressRejectInvalidAPIKey, rejectReason) } func TestApiKeyAuthWithSubscriptionGoogle_MarksUnavailableGroupBusinessLimited(t *testing.T) { @@ -422,9 +485,12 @@ func TestApiKeyAuthWithSubscriptionGoogle_MarksUnavailableGroupBusinessLimited(t r := gin.New() var markedBusinessLimited bool var businessLimitedReason string + var rejectReason IngressRejectReason + var rejected bool r.Use(func(c *gin.Context) { c.Next() markedBusinessLimited = service.HasOpsClientBusinessLimited(c) + rejectReason, rejected = GetIngressRejectReason(c) if v, ok := c.Get(service.OpsClientBusinessLimitedReasonKey); ok { businessLimitedReason, _ = v.(string) } @@ -452,6 +518,8 @@ func TestApiKeyAuthWithSubscriptionGoogle_MarksUnavailableGroupBusinessLimited(t require.Equal(t, "API Key 所属分组已删除", resp.Error.Message) require.True(t, markedBusinessLimited) require.Equal(t, service.OpsClientBusinessLimitedReasonAPIKeyGroupUnavailable, businessLimitedReason) + require.True(t, rejected) + require.Equal(t, IngressRejectGroupDeleted, rejectReason) } func TestApiKeyAuthWithSubscriptionGoogle_RepoError(t *testing.T) { diff --git a/backend/internal/server/middleware/api_key_auth_test.go b/backend/internal/server/middleware/api_key_auth_test.go index 2a2e2f558..a3216f8c2 100644 --- a/backend/internal/server/middleware/api_key_auth_test.go +++ b/backend/internal/server/middleware/api_key_auth_test.go @@ -8,6 +8,8 @@ import ( "errors" "net/http" "net/http/httptest" + "strings" + "sync/atomic" "testing" "time" @@ -19,6 +21,35 @@ import ( "github.com/stretchr/testify/require" ) +func TestAPIKeyAuthRejectsOversizedCredentialsBeforeLookup(t *testing.T) { + gin.SetMode(gin.TestMode) + var calls atomic.Int32 + repo := &stubApiKeyRepo{getByKey: func(context.Context, string) (*service.APIKey, error) { + calls.Add(1) + return nil, service.ErrAPIKeyNotFound + }} + cfg := &config.Config{RunMode: config.RunModeSimple} + svc := service.NewAPIKeyService(repo, nil, nil, nil, nil, nil, cfg) + + for _, headers := range []map[string]string{ + {"x-api-key": strings.Repeat("x", service.MaxAPIKeyCredentialBytes+1)}, + {"Authorization": "Bearer " + strings.Repeat("x", service.MaxAPIKeyCredentialBytes+1)}, + {"Authorization": strings.Repeat("x", maxAPIKeyAuthorizationHeaderBytes+1)}, + } { + r := gin.New() + r.Use(gin.HandlerFunc(NewAPIKeyAuthMiddleware(svc, nil, cfg))) + r.GET("/t", func(c *gin.Context) { c.Status(http.StatusOK) }) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/t", nil) + for name, value := range headers { + req.Header.Set(name, value) + } + r.ServeHTTP(w, req) + require.Equal(t, http.StatusUnauthorized, w.Code) + } + require.Zero(t, calls.Load()) +} + func TestSimpleModeBypassesQuotaCheck(t *testing.T) { gin.SetMode(gin.TestMode) @@ -436,6 +467,7 @@ func TestAPIKeyAuthRejectsUnavailableGroup(t *testing.T) { wantStatus int wantCode string wantMarked bool + wantReject IngressRejectReason }{ { name: "active group passes", @@ -460,6 +492,7 @@ func TestAPIKeyAuthRejectsUnavailableGroup(t *testing.T) { wantStatus: http.StatusForbidden, wantCode: "GROUP_DISABLED", wantMarked: true, + wantReject: IngressRejectGroupDisabled, }, { name: "deleted status group is forbidden", @@ -473,6 +506,7 @@ func TestAPIKeyAuthRejectsUnavailableGroup(t *testing.T) { wantStatus: http.StatusForbidden, wantCode: "GROUP_DELETED", wantMarked: true, + wantReject: IngressRejectGroupDeleted, }, { name: "missing group edge is forbidden", @@ -480,6 +514,7 @@ func TestAPIKeyAuthRejectsUnavailableGroup(t *testing.T) { wantStatus: http.StatusForbidden, wantCode: "GROUP_DELETED", wantMarked: true, + wantReject: IngressRejectGroupDeleted, }, } @@ -508,9 +543,12 @@ func TestAPIKeyAuthRejectsUnavailableGroup(t *testing.T) { router := gin.New() var markedBusinessLimited bool var businessLimitedReason string + var rejectReason IngressRejectReason + var rejected bool router.Use(func(c *gin.Context) { c.Next() markedBusinessLimited = service.HasOpsClientBusinessLimited(c) + rejectReason, rejected = GetIngressRejectReason(c) if v, ok := c.Get(service.OpsClientBusinessLimitedReasonKey); ok { businessLimitedReason, _ = v.(string) } @@ -530,6 +568,8 @@ func TestAPIKeyAuthRejectsUnavailableGroup(t *testing.T) { require.Contains(t, w.Body.String(), tt.wantCode) } require.Equal(t, tt.wantMarked, markedBusinessLimited) + require.Equal(t, tt.wantReject != "", rejected) + require.Equal(t, tt.wantReject, rejectReason) if tt.wantMarked { require.Equal(t, service.OpsClientBusinessLimitedReasonAPIKeyGroupUnavailable, businessLimitedReason) } @@ -537,6 +577,112 @@ func TestAPIKeyAuthRejectsUnavailableGroup(t *testing.T) { } } +func TestAPIKeyAuthMarksOnlyExpectedIngressRejections(t *testing.T) { + gin.SetMode(gin.TestMode) + + tests := []struct { + name string + path string + key string + authHeader string + repoErr error + wantStatus int + wantCode string + wantReason IngressRejectReason + }{ + { + name: "query key deprecated", + path: "/t?key=legacy", + wantStatus: http.StatusBadRequest, + wantCode: "api_key_in_query_deprecated", + wantReason: IngressRejectQueryAPIKeyDeprecated, + }, + { + name: "missing key", + path: "/t", + wantStatus: http.StatusUnauthorized, + wantCode: "API_KEY_REQUIRED", + wantReason: IngressRejectAPIKeyRequired, + }, + { + name: "malformed authorization", + path: "/t", + authHeader: "Basic not-a-bearer-key", + wantStatus: http.StatusUnauthorized, + wantCode: "API_KEY_REQUIRED", + wantReason: IngressRejectInvalidAPIKey, + }, + { + name: "oversized key", + path: "/t", + key: strings.Repeat("x", service.MaxAPIKeyCredentialBytes+1), + wantStatus: http.StatusUnauthorized, + wantCode: "INVALID_API_KEY", + wantReason: IngressRejectInvalidAPIKey, + }, + { + name: "invalid key", + path: "/t", + key: "invalid", + repoErr: service.ErrAPIKeyNotFound, + wantStatus: http.StatusUnauthorized, + wantCode: "INVALID_API_KEY", + wantReason: IngressRejectInvalidAPIKey, + }, + { + name: "repository failure remains operational error", + path: "/t", + key: "valid-shape", + repoErr: errors.New("database unavailable"), + wantStatus: http.StatusInternalServerError, + wantCode: "INTERNAL_ERROR", + }, + { + name: "auth lookup bulkhead rejection is an admission rejection", + path: "/t", + key: "valid-shape", + repoErr: service.ErrAPIKeyAuthOverloaded, + wantStatus: http.StatusServiceUnavailable, + wantCode: "API_KEY_AUTH_OVERLOADED", + wantReason: IngressRejectAPIKeyAuthOverloaded, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + repo := &stubApiKeyRepo{getByKey: func(context.Context, string) (*service.APIKey, error) { + return nil, tt.repoErr + }} + cfg := &config.Config{RunMode: config.RunModeSimple} + apiKeyService := service.NewAPIKeyService(repo, nil, nil, nil, nil, nil, cfg) + router := gin.New() + var reason IngressRejectReason + var rejected bool + router.Use(func(c *gin.Context) { + c.Next() + reason, rejected = GetIngressRejectReason(c) + }) + router.Use(gin.HandlerFunc(NewAPIKeyAuthMiddleware(apiKeyService, nil, cfg))) + router.GET("/t", func(c *gin.Context) { c.Status(http.StatusOK) }) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, tt.path, nil) + if tt.key != "" { + req.Header.Set("x-api-key", tt.key) + } + if tt.authHeader != "" { + req.Header.Set("Authorization", tt.authHeader) + } + router.ServeHTTP(w, req) + + require.Equal(t, tt.wantStatus, w.Code) + require.Contains(t, w.Body.String(), tt.wantCode) + require.Equal(t, tt.wantReason != "", rejected) + require.Equal(t, tt.wantReason, reason) + }) + } +} + func TestAPIKeyAuthSetsOpsFallbackKeyOnEarlyAbort(t *testing.T) { gin.SetMode(gin.TestMode) @@ -686,9 +832,12 @@ func TestRequireGroupAssignmentMarksUngroupedKeyBusinessLimited(t *testing.T) { router := gin.New() var markedBusinessLimited bool var businessLimitedReason string + var rejectReason IngressRejectReason + var rejected bool router.Use(func(c *gin.Context) { c.Next() markedBusinessLimited = service.HasOpsClientBusinessLimited(c) + rejectReason, rejected = GetIngressRejectReason(c) if v, ok := c.Get(service.OpsClientBusinessLimitedReasonKey); ok { businessLimitedReason, _ = v.(string) } @@ -708,6 +857,8 @@ func TestRequireGroupAssignmentMarksUngroupedKeyBusinessLimited(t *testing.T) { require.Equal(t, http.StatusForbidden, w.Code) require.Contains(t, w.Body.String(), "not assigned to any group") + require.True(t, rejected) + require.Equal(t, IngressRejectGroupUnassigned, rejectReason) require.True(t, markedBusinessLimited) require.Equal(t, service.OpsClientBusinessLimitedReasonAPIKeyGroupUnassigned, businessLimitedReason) } @@ -822,7 +973,7 @@ func TestAPIKeyAuthIPRestrictionIncludesClientIPForBlacklistDenial(t *testing.T) requireAPIKeyAuthError(t, w, "ACCESS_DENIED", "Access denied. Your IP is 9.9.9.9") } -func TestAPIKeyAuthIPRestrictionCanTrustForwardedClientIPForReverseProxy(t *testing.T) { +func TestAPIKeyAuthIPRestrictionUsesConfiguredTrustedProxy(t *testing.T) { gin.SetMode(gin.TestMode) user := &service.User{ @@ -855,7 +1006,7 @@ func TestAPIKeyAuthIPRestrictionCanTrustForwardedClientIPForReverseProxy(t *test cfg.SetTrustForwardedIPForAPIKeyACL(true) apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg) router := gin.New() - require.NoError(t, router.SetTrustedProxies(nil)) + require.NoError(t, router.SetTrustedProxies([]string{"9.9.9.9"})) router.Use(gin.HandlerFunc(NewAPIKeyAuthMiddleware(apiKeyService, nil, cfg))) router.GET("/t", func(c *gin.Context) { c.JSON(http.StatusOK, gin.H{"ok": true}) @@ -906,7 +1057,7 @@ func TestAPIKeyAuthIPRestrictionUsesForwardedClientIPInDenialWhenTrusted(t *test cfg.SetTrustForwardedIPForAPIKeyACL(true) apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg) router := gin.New() - require.NoError(t, router.SetTrustedProxies(nil)) + require.NoError(t, router.SetTrustedProxies([]string{"9.9.9.9"})) router.Use(gin.HandlerFunc(NewAPIKeyAuthMiddleware(apiKeyService, nil, cfg))) router.GET("/t", func(c *gin.Context) { c.JSON(http.StatusOK, gin.H{"ok": true}) diff --git a/backend/internal/server/middleware/client_request_id.go b/backend/internal/server/middleware/client_request_id.go index 5f886646e..7bf1c8af0 100644 --- a/backend/internal/server/middleware/client_request_id.go +++ b/backend/internal/server/middleware/client_request_id.go @@ -24,7 +24,14 @@ func ClientRequestID() gin.HandlerFunc { } if v, _ := c.Request.Context().Value(ctxkey.ClientRequestID).(string); strings.TrimSpace(v) != "" { - c.Header(clientRequestIDHeader, strings.TrimSpace(v)) + var valid bool + v, valid = normalizeCorrelationID(v) + if !valid { + v = uuid.New().String() + } + c.Header(clientRequestIDHeader, v) + ctx := context.WithValue(c.Request.Context(), ctxkey.ClientRequestID, v) + c.Request = c.Request.WithContext(ctx) c.Next() return } diff --git a/backend/internal/server/middleware/client_request_id_test.go b/backend/internal/server/middleware/client_request_id_test.go index 394c16125..c57c91fad 100644 --- a/backend/internal/server/middleware/client_request_id_test.go +++ b/backend/internal/server/middleware/client_request_id_test.go @@ -4,6 +4,7 @@ import ( "context" "net/http" "net/http/httptest" + "strings" "testing" "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" @@ -30,6 +31,23 @@ func TestClientRequestIDGeneratesAndExposesID(t *testing.T) { require.Equal(t, w.Body.String(), w.Header().Get(clientRequestIDHeader)) } +func TestClientRequestIDBoundsExistingContextID(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(ClientRequestID()) + router.GET("/", func(c *gin.Context) { + value, _ := c.Request.Context().Value(ctxkey.ClientRequestID).(string) + c.String(http.StatusOK, value) + }) + req := httptest.NewRequest(http.MethodGet, "/", nil) + req = req.WithContext(context.WithValue(req.Context(), ctxkey.ClientRequestID, strings.Repeat("x", 200))) + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + require.Len(t, w.Body.String(), 36) + require.NotEqual(t, strings.Repeat("x", maxPersistentRequestIDBytes), w.Body.String()) + require.Equal(t, w.Body.String(), w.Header().Get(clientRequestIDHeader)) +} + func TestClientRequestIDPreservesExistingContextID(t *testing.T) { gin.SetMode(gin.TestMode) router := gin.New() diff --git a/backend/internal/server/middleware/ingress_reject.go b/backend/internal/server/middleware/ingress_reject.go new file mode 100644 index 000000000..e61350d78 --- /dev/null +++ b/backend/internal/server/middleware/ingress_reject.go @@ -0,0 +1,166 @@ +package middleware + +import ( + "math" + "net/netip" + "strconv" + "strings" + "sync/atomic" + "time" + + "github.com/gin-gonic/gin" +) + +// IngressRejectReason identifies expected gateway admission failures that must +// not be treated as operational request errors. +type IngressRejectReason string + +const ( + IngressRejectQueryAPIKeyDeprecated IngressRejectReason = "query_api_key_deprecated" + IngressRejectAPIKeyRequired IngressRejectReason = "api_key_required" + IngressRejectInvalidAPIKey IngressRejectReason = "invalid_api_key" + IngressRejectAPIKeyDisabled IngressRejectReason = "api_key_disabled" + IngressRejectIPRestricted IngressRejectReason = "ip_restricted" + IngressRejectUserInactive IngressRejectReason = "user_inactive" + IngressRejectGroupDeleted IngressRejectReason = "group_deleted" + IngressRejectGroupDisabled IngressRejectReason = "group_disabled" + IngressRejectGroupNotAllowed IngressRejectReason = "group_not_allowed" + IngressRejectGroupUnassigned IngressRejectReason = "group_unassigned" + IngressRejectInvalidAuthRateLimited IngressRejectReason = "invalid_auth_rate_limited" + IngressRejectAPIKeyAuthOverloaded IngressRejectReason = "api_key_auth_overloaded" +) + +const ingressRejectReasonContextKey = "ingress_reject_reason" + +type IngressRejectRecorder interface { + RecordIngressReject(reason, routeFamily, protocol, clientIP string, userID, apiKeyID int64) +} + +func invalidAuthClientKey(c *gin.Context) string { + return normalizeIngressRejectIP(SecurityClientIP(c)) +} + +func rejectInvalidAuthAbuse(c *gin.Context, apiKeyService interface { + CheckInvalidAuthAbuse(string) (time.Duration, bool) +}) bool { + if c == nil || apiKeyService == nil { + return false + } + retry, blocked := apiKeyService.CheckInvalidAuthAbuse(invalidAuthClientKey(c)) + if !blocked { + return false + } + retrySeconds := int(math.Ceil(retry.Seconds())) + if retrySeconds < 1 { + retrySeconds = 1 + } + c.Header("Retry-After", strconv.Itoa(retrySeconds)) + MarkIngressRejected(c, IngressRejectInvalidAuthRateLimited) + return true +} + +func recordInvalidAuthFailure(c *gin.Context, apiKeyService interface { + RecordInvalidAuthFailure(string) +}) { + if c == nil || apiKeyService == nil { + return + } + apiKeyService.RecordInvalidAuthFailure(invalidAuthClientKey(c)) +} + +type ingressRejectRecorderHolder struct{ recorder IngressRejectRecorder } + +var activeIngressRejectRecorder atomic.Pointer[ingressRejectRecorderHolder] + +func SetIngressRejectRecorder(recorder IngressRejectRecorder) { + if recorder == nil { + activeIngressRejectRecorder.Store(nil) + return + } + activeIngressRejectRecorder.Store(&ingressRejectRecorderHolder{recorder: recorder}) +} + +// MarkIngressRejected marks a request as rejected before gateway admission. +func MarkIngressRejected(c *gin.Context, reason IngressRejectReason) { + if c == nil || reason == "" { + return + } + c.Set(ingressRejectReasonContextKey, reason) +} + +// GetIngressRejectReason returns the admission rejection reason, if any. +func GetIngressRejectReason(c *gin.Context) (IngressRejectReason, bool) { + if c == nil { + return "", false + } + value, exists := c.Get(ingressRejectReasonContextKey) + if !exists { + return "", false + } + reason, ok := value.(IngressRejectReason) + return reason, ok && reason != "" +} + +func recordIngressReject(c *gin.Context, reason IngressRejectReason) { + holder := activeIngressRejectRecorder.Load() + if holder == nil || holder.recorder == nil || c == nil || c.Request == nil { + return + } + routeFamily, protocol := ingressRejectRoute(c.Request.URL.Path) + clientIP := normalizeIngressRejectIP(SecurityClientIP(c)) + var userID, apiKeyID int64 + if apiKey, ok := GetAPIKeyFromContext(c); ok && apiKey != nil { + apiKeyID = apiKey.ID + if apiKey.User != nil { + userID = apiKey.User.ID + } + } else if apiKey, ok := GetOpsFallbackAPIKey(c); ok && apiKey != nil { + apiKeyID = apiKey.ID + if apiKey.User != nil { + userID = apiKey.User.ID + } + } + holder.recorder.RecordIngressReject(string(reason), routeFamily, protocol, clientIP, userID, apiKeyID) +} + +func normalizeIngressRejectIP(raw string) string { + addr, err := netip.ParseAddr(strings.TrimSpace(raw)) + if err != nil { + return "0.0.0.0" + } + addr = addr.Unmap() + if addr.Is6() { + return netip.PrefixFrom(addr, 64).Masked().Addr().String() + } + return addr.String() +} + +func ingressRejectRoute(path string) (string, string) { + path = strings.ToLower(strings.TrimSpace(path)) + switch { + case strings.HasPrefix(path, "/antigravity/v1beta"): + return "antigravity", "google" + case strings.HasPrefix(path, "/v1beta"): + return "gemini", "google" + case strings.HasPrefix(path, "/backend-api/codex"): + return "codex", "openai" + case strings.HasPrefix(path, "/antigravity"): + return "antigravity", "anthropic" + case strings.Contains(path, "/messages"): + return "messages", "anthropic" + case strings.Contains(path, "/responses"): + return "responses", "openai" + case strings.Contains(path, "/chat/completions"): + return "chat_completions", "openai" + case strings.Contains(path, "/images"): + return "images", "openai" + case strings.Contains(path, "/videos"): + return "videos", "openai" + case strings.Contains(path, "/embeddings"): + return "embeddings", "openai" + case strings.Contains(path, "/models"): + return "models", "openai" + default: + return "other", "gateway" + } +} diff --git a/backend/internal/server/middleware/ingress_reject_access_sampler.go b/backend/internal/server/middleware/ingress_reject_access_sampler.go new file mode 100644 index 000000000..98acb7893 --- /dev/null +++ b/backend/internal/server/middleware/ingress_reject_access_sampler.go @@ -0,0 +1,58 @@ +package middleware + +import ( + "sync" + "time" +) + +const ( + ingressRejectAccessLogLimit = 20 + ingressRejectAccessLogWindow = time.Second + ingressRejectDroppedSummaryPeriod = 30 * time.Second +) + +type ingressRejectAccessSampler struct { + mu sync.Mutex + limit int + window time.Duration + summaryPeriod time.Duration + windowStart time.Time + emitted int + dropped uint64 + lastSummary time.Time +} + +func newIngressRejectAccessSampler(limit int, window, summaryPeriod time.Duration) *ingressRejectAccessSampler { + return &ingressRejectAccessSampler{limit: limit, window: window, summaryPeriod: summaryPeriod} +} + +// allow applies one process-wide fixed-window budget. It stores no attacker +// dimensions, so memory remains constant even for rotating keys and addresses. +func (s *ingressRejectAccessSampler) allow(now time.Time) (allowed bool, droppedSummary uint64) { + if s == nil || s.limit <= 0 || s.window <= 0 { + return false, 0 + } + s.mu.Lock() + defer s.mu.Unlock() + if s.windowStart.IsZero() || now.Sub(s.windowStart) >= s.window || now.Before(s.windowStart) { + s.windowStart = now + s.emitted = 0 + } + if s.emitted < s.limit { + s.emitted++ + return true, 0 + } + s.dropped++ + if s.summaryPeriod > 0 && (s.lastSummary.IsZero() || now.Sub(s.lastSummary) >= s.summaryPeriod) { + droppedSummary = s.dropped + s.dropped = 0 + s.lastSummary = now + } + return false, droppedSummary +} + +var globalIngressRejectAccessSampler = newIngressRejectAccessSampler( + ingressRejectAccessLogLimit, + ingressRejectAccessLogWindow, + ingressRejectDroppedSummaryPeriod, +) diff --git a/backend/internal/server/middleware/ingress_reject_access_sampler_test.go b/backend/internal/server/middleware/ingress_reject_access_sampler_test.go new file mode 100644 index 000000000..5ca11f133 --- /dev/null +++ b/backend/internal/server/middleware/ingress_reject_access_sampler_test.go @@ -0,0 +1,81 @@ +package middleware + +import ( + "net/http" + "net/http/httptest" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestIngressRejectAccessSamplerConcurrentGlobalLimit(t *testing.T) { + sampler := newIngressRejectAccessSampler(10, time.Hour, time.Minute) + now := time.Now() + var allowed atomic.Int64 + var wg sync.WaitGroup + for i := 0; i < 200; i++ { + wg.Add(1) + go func() { + defer wg.Done() + if ok, _ := sampler.allow(now); ok { + allowed.Add(1) + } + }() + } + wg.Wait() + require.Equal(t, int64(10), allowed.Load()) +} + +func TestLoggerIngressRejectSamplingIsBoundedAndSummarySkipsOpsSink(t *testing.T) { + gin.SetMode(gin.TestMode) + original := globalIngressRejectAccessSampler + globalIngressRejectAccessSampler = newIngressRejectAccessSampler(2, time.Hour, time.Hour) + t.Cleanup(func() { globalIngressRejectAccessSampler = original }) + sink := initMiddlewareTestLogger(t) + router := gin.New() + router.Use(Logger()) + router.GET("/v1/messages", func(c *gin.Context) { + MarkIngressRejected(c, IngressRejectInvalidAPIKey) + c.Status(http.StatusUnauthorized) + }) + for i := 0; i < 20; i++ { + router.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, "/v1/messages", nil)) + } + var accessEvents, summaries int + for _, event := range sink.list() { + switch event.Message { + case "http request completed": + accessEvents++ + case "ingress rejection access logs dropped": + summaries++ + if skipped, _ := event.Fields[logger.OpsSystemLogSkipField].(bool); !skipped { + t.Fatalf("dropped summary must skip ops system log sink") + } + } + } + require.Equal(t, 2, accessEvents) + require.Equal(t, 1, summaries) +} + +func TestIngressRejectAccessSamplerDroppedSummaryIsLowFrequency(t *testing.T) { + sampler := newIngressRejectAccessSampler(1, time.Hour, time.Second) + now := time.Now() + allowed, summary := sampler.allow(now) + require.True(t, allowed) + require.Zero(t, summary) + + allowed, summary = sampler.allow(now.Add(100 * time.Millisecond)) + require.False(t, allowed) + require.Equal(t, uint64(1), summary) + allowed, summary = sampler.allow(now.Add(200 * time.Millisecond)) + require.False(t, allowed) + require.Zero(t, summary) + allowed, summary = sampler.allow(now.Add(2 * time.Second)) + require.False(t, allowed) + require.Equal(t, uint64(2), summary) +} diff --git a/backend/internal/server/middleware/ingress_reject_test.go b/backend/internal/server/middleware/ingress_reject_test.go new file mode 100644 index 000000000..fec19a854 --- /dev/null +++ b/backend/internal/server/middleware/ingress_reject_test.go @@ -0,0 +1,50 @@ +package middleware + +import ( + "net/http" + "net/http/httptest" + "sync" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +type ingressRejectRecorderStub struct { + mu sync.Mutex + calls int + clientIP string +} + +func (r *ingressRejectRecorderStub) RecordIngressReject(_, _, _, clientIP string, _, _ int64) { + r.mu.Lock() + defer r.mu.Unlock() + r.calls++ + r.clientIP = clientIP +} + +func TestNormalizeIngressRejectIP(t *testing.T) { + require.Equal(t, "2001:db8:abcd:1234::", normalizeIngressRejectIP("2001:db8:abcd:1234:ffff::1")) + require.Equal(t, "192.0.2.4", normalizeIngressRejectIP("::ffff:192.0.2.4")) + require.Equal(t, "0.0.0.0", normalizeIngressRejectIP("not-an-ip")) +} + +func TestLoggerRecordsIngressRejectOnce(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := &ingressRejectRecorderStub{} + SetIngressRejectRecorder(recorder) + t.Cleanup(func() { SetIngressRejectRecorder(nil) }) + router := gin.New() + router.Use(Logger()) + router.GET("/v1/messages", func(c *gin.Context) { + MarkIngressRejected(c, IngressRejectInvalidAPIKey) + c.Status(http.StatusUnauthorized) + }) + request := httptest.NewRequest(http.MethodGet, "/v1/messages", nil) + request.RemoteAddr = "[2001:db8:abcd:1234:ffff::1]:1234" + router.ServeHTTP(httptest.NewRecorder(), request) + recorder.mu.Lock() + require.Equal(t, 1, recorder.calls) + require.Equal(t, "2001:db8:abcd:1234::", recorder.clientIP) + recorder.mu.Unlock() +} diff --git a/backend/internal/server/middleware/invalid_auth_abuse_test.go b/backend/internal/server/middleware/invalid_auth_abuse_test.go new file mode 100644 index 000000000..f8242146b --- /dev/null +++ b/backend/internal/server/middleware/invalid_auth_abuse_test.go @@ -0,0 +1,147 @@ +//go:build unit + +package middleware + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func invalidAuthAbuseTestConfig(threshold int) *config.Config { + return &config.Config{ + RunMode: config.RunModeSimple, + APIKeyAuth: config.APIKeyAuthCacheConfig{InvalidAbuse: config.InvalidAuthAbuseConfig{ + Enabled: true, Threshold: threshold, WindowSeconds: 60, BlockSeconds: 60, Capacity: 256, + }}, + } +} + +func TestAPIKeyAuthInvalidAbuseReturns429BeforeRepository(t *testing.T) { + gin.SetMode(gin.TestMode) + repoCalls := 0 + repo := &stubApiKeyRepo{getByKey: func(context.Context, string) (*service.APIKey, error) { + repoCalls++ + return nil, service.ErrAPIKeyNotFound + }} + cfg := invalidAuthAbuseTestConfig(3) + svc := service.NewAPIKeyService(repo, nil, nil, nil, nil, nil, cfg) + r := gin.New() + var reason IngressRejectReason + r.Use(func(c *gin.Context) { c.Next(); reason, _ = GetIngressRejectReason(c) }) + r.Use(gin.HandlerFunc(NewAPIKeyAuthMiddleware(svc, nil, cfg))) + r.POST("/v1/messages", func(c *gin.Context) { c.Status(http.StatusOK) }) + + requests := []*http.Request{ + httpRequest(t, "/v1/messages", "", ""), + httpRequest(t, "/v1/messages", "Basic malformed", ""), + httpRequest(t, "/v1/messages", "", "random-invalid-key"), + } + for _, req := range requests { + w := httptest.NewRecorder() + r.ServeHTTP(w, req) + require.NotEqual(t, http.StatusTooManyRequests, w.Code) + } + + w := httptest.NewRecorder() + r.ServeHTTP(w, httpRequest(t, "/v1/messages", "", "another-random-key")) + require.Equal(t, http.StatusTooManyRequests, w.Code) + require.Equal(t, "60", w.Header().Get("Retry-After")) + require.Contains(t, w.Body.String(), "INVALID_AUTH_RATE_LIMITED") + require.Equal(t, IngressRejectInvalidAuthRateLimited, reason) + require.Equal(t, 1, repoCalls, "rate-limited request must not reach the repository") +} + +func TestGoogleAPIKeyAuthInvalidAbuseReturnsProtocol429(t *testing.T) { + gin.SetMode(gin.TestMode) + repoCalls := 0 + repo := fakeAPIKeyRepo{getByKey: func(context.Context, string) (*service.APIKey, error) { + repoCalls++ + return nil, service.ErrAPIKeyNotFound + }} + cfg := invalidAuthAbuseTestConfig(2) + svc := service.NewAPIKeyService(repo, nil, nil, nil, nil, nil, cfg) + r := gin.New() + var reason IngressRejectReason + r.Use(func(c *gin.Context) { c.Next(); reason, _ = GetIngressRejectReason(c) }) + r.Use(APIKeyAuthGoogle(svc, cfg)) + r.POST("/v1beta/models/test:generateContent", func(c *gin.Context) { c.Status(http.StatusOK) }) + for _, key := range []string{"random-1", "random-2"} { + w := httptest.NewRecorder() + req := httpRequest(t, "/v1beta/models/test:generateContent", "", key) + req.Header.Del("x-api-key") + req.Header.Set("x-goog-api-key", key) + r.ServeHTTP(w, req) + require.Equal(t, http.StatusUnauthorized, w.Code) + } + w := httptest.NewRecorder() + req := httpRequest(t, "/v1beta/models/test:generateContent", "", "random-3") + req.Header.Del("x-api-key") + req.Header.Set("x-goog-api-key", "random-3") + r.ServeHTTP(w, req) + require.Equal(t, http.StatusTooManyRequests, w.Code) + require.Equal(t, "60", w.Header().Get("Retry-After")) + require.Contains(t, w.Body.String(), "RESOURCE_EXHAUSTED") + require.Equal(t, IngressRejectInvalidAuthRateLimited, reason) + require.Equal(t, 2, repoCalls) +} + +func TestInvalidAuthAbuseDoesNotCountValidOrOperationalFailures(t *testing.T) { + gin.SetMode(gin.TestMode) + user := &service.User{ID: 1, Status: service.StatusActive, Role: service.RoleUser, Balance: 1} + repo := &stubApiKeyRepo{getByKey: func(_ context.Context, key string) (*service.APIKey, error) { + switch key { + case "valid-key": + return &service.APIKey{ID: 1, UserID: 1, Key: key, Status: service.StatusActive, User: user}, nil + case "db-error": + return nil, errors.New("database unavailable") + default: + return nil, service.ErrAPIKeyNotFound + } + }} + cfg := invalidAuthAbuseTestConfig(10) + svc := service.NewAPIKeyService(repo, nil, nil, nil, nil, nil, cfg) + r := gin.New() + r.Use(gin.HandlerFunc(NewAPIKeyAuthMiddleware(svc, nil, cfg))) + r.POST("/t", func(c *gin.Context) { c.Status(http.StatusOK) }) + + for _, tc := range []struct { + key string + want int + }{{"invalid", 401}, {"valid-key", 200}, {"db-error", 500}, {"db-error", 500}} { + w := httptest.NewRecorder() + r.ServeHTTP(w, httpRequest(t, "/t", "", tc.key)) + require.Equal(t, tc.want, w.Code) + } + w := httptest.NewRecorder() + req := httpRequest(t, "/t", "", "") + req.Header.Set("x-goog-api-key", "valid-key") + r.ServeHTTP(w, req) + require.Equal(t, http.StatusOK, w.Code) + require.Equal(t, uint64(1), svc.InvalidAuthAbuseHealth().Recorded) +} + +func TestNormalizeIngressRejectIPGroupsIPv6By64(t *testing.T) { + require.Equal(t, "2001:db8:abcd:1234::", normalizeIngressRejectIP("2001:db8:abcd:1234:1111::1")) + require.Equal(t, normalizeIngressRejectIP("2001:db8:abcd:1234:1111::1"), normalizeIngressRejectIP("2001:db8:abcd:1234:ffff::2")) +} + +func httpRequest(t *testing.T, path, authorization, apiKey string) *http.Request { + t.Helper() + req := httptest.NewRequest(http.MethodPost, path, nil) + req.RemoteAddr = "203.0.113.10:12345" + if authorization != "" { + req.Header.Set("Authorization", authorization) + } + if apiKey != "" { + req.Header.Set("x-api-key", apiKey) + } + return req +} diff --git a/backend/internal/server/middleware/logger.go b/backend/internal/server/middleware/logger.go index 6cec1c1a0..a9189bb05 100644 --- a/backend/internal/server/middleware/logger.go +++ b/backend/internal/server/middleware/logger.go @@ -37,6 +37,21 @@ func Logger() gin.HandlerFunc { accountID, hasAccountID := c.Request.Context().Value(ctxkey.AccountID).(int64) platform, _ := c.Request.Context().Value(ctxkey.Platform).(string) model, _ := c.Request.Context().Value(ctxkey.Model).(string) + reason, rejected := GetIngressRejectReason(c) + if rejected { + recordIngressReject(c, reason) + allowed, droppedSummary := globalIngressRejectAccessSampler.allow(endTime) + if droppedSummary > 0 { + logger.FromContext(c.Request.Context()).Info("ingress rejection access logs dropped", + zap.String("component", "http.access"), + zap.Uint64("dropped_count", droppedSummary), + zap.Bool(logger.OpsSystemLogSkipField, true), + ) + } + if !allowed { + return + } + } fields := []zap.Field{ zap.String("component", "http.access"), @@ -47,6 +62,12 @@ func Logger() gin.HandlerFunc { zap.String("method", method), zap.String("path", path), } + if rejected { + fields = append(fields, + zap.String("ingress_reject_reason", string(reason)), + zap.Bool(logger.OpsSystemLogSkipField, true), + ) + } if hasAccountID && accountID > 0 { fields = append(fields, zap.Int64("account_id", accountID)) } diff --git a/backend/internal/server/middleware/middleware.go b/backend/internal/server/middleware/middleware.go index 9efe78a33..0556be2a5 100644 --- a/backend/internal/server/middleware/middleware.go +++ b/backend/internal/server/middleware/middleware.go @@ -121,6 +121,7 @@ func RequireGroupAssignment(settingService *service.SettingService, writeError G return } service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonAPIKeyGroupUnassigned) + MarkIngressRejected(c, IngressRejectGroupUnassigned) writeError(c, http.StatusForbidden, "API Key is not assigned to any group and cannot be used. Please contact the administrator to assign it to a group.") c.Abort() } diff --git a/backend/internal/server/middleware/request_access_logger_test.go b/backend/internal/server/middleware/request_access_logger_test.go index bf3666f2e..ded191964 100644 --- a/backend/internal/server/middleware/request_access_logger_test.go +++ b/backend/internal/server/middleware/request_access_logger_test.go @@ -112,6 +112,26 @@ func TestRequestLogger_KeepIncomingRequestID(t *testing.T) { } } +func TestRequestLoggerBoundsIncomingRequestID(t *testing.T) { + gin.SetMode(gin.TestMode) + r := gin.New() + r.Use(RequestLogger()) + r.GET("/t", func(c *gin.Context) { + reqID, _ := c.Request.Context().Value(ctxkey.RequestID).(string) + if len(reqID) != 36 { + t.Fatalf("request_id length=%d", len(reqID)) + } + c.Status(http.StatusOK) + }) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/t", nil) + req.Header.Set(requestIDHeader, strings.Repeat("r", 1024)) + r.ServeHTTP(w, req) + if got := len(w.Header().Get(requestIDHeader)); got != 36 { + t.Fatalf("response request_id length=%d", got) + } +} + func TestLogger_AccessLogIncludesCoreFields(t *testing.T) { gin.SetMode(gin.TestMode) sink := initMiddlewareTestLogger(t) @@ -180,11 +200,43 @@ func TestLogger_AccessLogIncludesCoreFields(t *testing.T) { } } -func TestLogger_AccessLogUsesForwardedClientIP(t *testing.T) { +func TestLogger_IngressRejectRemainsInStandardAccessLog(t *testing.T) { + gin.SetMode(gin.TestMode) + sink := initMiddlewareTestLogger(t) + r := gin.New() + r.Use(Logger()) + r.GET("/v1/messages", func(c *gin.Context) { + MarkIngressRejected(c, IngressRejectInvalidAPIKey) + c.Status(http.StatusUnauthorized) + }) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/v1/messages", nil) + r.ServeHTTP(w, req) + + if w.Code != http.StatusUnauthorized { + t.Fatalf("status=%d", w.Code) + } + events := sink.list() + if len(events) != 1 { + t.Fatalf("events=%d, want 1", len(events)) + } + if got := events[0].Fields["ingress_reject_reason"]; got != string(IngressRejectInvalidAPIKey) { + t.Fatalf("ingress_reject_reason=%v", got) + } + if got, _ := events[0].Fields[logger.OpsSystemLogSkipField].(bool); !got { + t.Fatalf("%s must be true", logger.OpsSystemLogSkipField) + } +} + +func TestLogger_AccessLogUsesForwardedClientIPFromTrustedProxy(t *testing.T) { gin.SetMode(gin.TestMode) sink := initMiddlewareTestLogger(t) r := gin.New() + if err := r.SetTrustedProxies([]string{"104.23.251.120"}); err != nil { + t.Fatalf("set trusted proxies: %v", err) + } r.Use(Logger()) r.GET("/api/test", func(c *gin.Context) { c.Status(http.StatusOK) @@ -193,7 +245,7 @@ func TestLogger_AccessLogUsesForwardedClientIP(t *testing.T) { w := httptest.NewRecorder() req := httptest.NewRequest(http.MethodGet, "/api/test", nil) req.RemoteAddr = "104.23.251.120:443" - req.Header.Set("CF-Connecting-IP", "203.0.113.42") + req.Header.Set("X-Forwarded-For", "203.0.113.42") r.ServeHTTP(w, req) if w.Code != http.StatusOK { t.Fatalf("status=%d", w.Code) diff --git a/backend/internal/server/middleware/request_logger.go b/backend/internal/server/middleware/request_logger.go index 0fb2feca1..e3d749744 100644 --- a/backend/internal/server/middleware/request_logger.go +++ b/backend/internal/server/middleware/request_logger.go @@ -21,14 +21,15 @@ func RequestLogger() gin.HandlerFunc { return } - requestID := strings.TrimSpace(c.GetHeader(requestIDHeader)) - if requestID == "" { + requestID, validRequestID := normalizeCorrelationID(c.GetHeader(requestIDHeader)) + if !validRequestID { requestID = uuid.NewString() } c.Header(requestIDHeader, requestID) ctx := context.WithValue(c.Request.Context(), ctxkey.RequestID, requestID) clientRequestID, _ := ctx.Value(ctxkey.ClientRequestID).(string) + clientRequestID, _ = normalizeCorrelationID(clientRequestID) requestLogger := logger.With( zap.String("component", "http"), diff --git a/backend/internal/server/middleware/request_metadata.go b/backend/internal/server/middleware/request_metadata.go new file mode 100644 index 000000000..740583cb1 --- /dev/null +++ b/backend/internal/server/middleware/request_metadata.go @@ -0,0 +1,30 @@ +package middleware + +import ( + "strings" + "unicode/utf8" +) + +const ( + maxPersistentRequestIDBytes = 64 + maxPersistentUserAgentBytes = 512 +) + +// normalizePersistentText bounds attacker-controlled metadata before it reaches +// logs or database columns while preserving valid UTF-8 content. +func normalizePersistentText(value string, maxBytes int) string { + value = strings.TrimSpace(strings.ToValidUTF8(value, "")) + if maxBytes <= 0 || len(value) <= maxBytes { + return value + } + value = value[:maxBytes] + for !utf8.ValidString(value) { + value = value[:len(value)-1] + } + return value +} + +func normalizeCorrelationID(value string) (string, bool) { + value = strings.TrimSpace(strings.ToValidUTF8(value, "")) + return value, value != "" && len(value) <= maxPersistentRequestIDBytes +} diff --git a/backend/internal/server/middleware/session_binding.go b/backend/internal/server/middleware/session_binding.go index 31edd110b..e535693c2 100644 --- a/backend/internal/server/middleware/session_binding.go +++ b/backend/internal/server/middleware/session_binding.go @@ -13,14 +13,15 @@ import ( // SessionBindingContext 全局中间件:将请求的客户端 IP 与 User-Agent 注入 // request context,供 token 签发路径(登录 / 刷新 / OAuth 回调)读取并写入会话绑定, // 同时作为审计日志、会话绑定校验的统一客户端 IP 来源。 -// IP 取值与 API Key IP 限制共用「信任反代传递的客户端 IP」系统开关: -// 开启时信任反代转发头(CF-Connecting-IP / X-Real-IP / X-Forwarded-For), -// 关闭时走 trusted_proxies 解析链,避免不可信头伪造绕过绑定。 +// IP 取值与 API Key IP 限制共用 Gin trusted_proxies 解析链;旧设置开关 +// 仅为配置兼容保留,不能单独使直连请求的转发头变为可信。 func SessionBindingContext(cfg *config.Config) gin.HandlerFunc { return func(c *gin.Context) { + userAgent := normalizePersistentText(c.Request.UserAgent(), maxPersistentUserAgentBytes) + c.Request.Header.Set("User-Agent", userAgent) binding := &service.SessionBinding{ IP: ip.GetSecurityClientIP(c, cfg.TrustForwardedIPForAPIKeyACL()), - UserAgent: c.Request.UserAgent(), + UserAgent: userAgent, } c.Request = c.Request.WithContext(service.WithSessionBinding(c.Request.Context(), binding)) c.Next() @@ -36,7 +37,7 @@ func requestSessionBinding(c *gin.Context) *service.SessionBinding { } return &service.SessionBinding{ IP: ip.GetTrustedClientIP(c), - UserAgent: c.Request.UserAgent(), + UserAgent: normalizePersistentText(c.Request.UserAgent(), maxPersistentUserAgentBytes), } } @@ -93,7 +94,7 @@ func enforceSessionBinding( Method: c.Request.Method, Path: path, ClientIP: binding.IP, - UserAgent: c.Request.UserAgent(), + UserAgent: normalizePersistentText(c.Request.UserAgent(), maxPersistentUserAgentBytes), StatusCode: 401, }) } diff --git a/backend/internal/server/middleware/session_binding_test.go b/backend/internal/server/middleware/session_binding_test.go index 2823dc6ee..839983623 100644 --- a/backend/internal/server/middleware/session_binding_test.go +++ b/backend/internal/server/middleware/session_binding_test.go @@ -4,6 +4,7 @@ package middleware import ( "net/http/httptest" + "strings" "testing" "github.com/Wei-Shaw/sub2api/internal/config" @@ -13,9 +14,7 @@ import ( "github.com/stretchr/testify/require" ) -// 反代场景:RemoteAddr 为 127.0.0.1,真实客户端 IP 在 X-Real-IP 中。 -// 会话绑定注入与审计 IP 必须与 API Key IP 限制共用「信任反代传递的客户端 IP」开关语义。 -func TestSessionBindingContextHonorsTrustForwardedToggle(t *testing.T) { +func TestSessionBindingContextDoesNotTrustHeadersWithoutTrustedProxy(t *testing.T) { gin.SetMode(gin.TestMode) for _, tc := range []struct { @@ -24,7 +23,7 @@ func TestSessionBindingContextHonorsTrustForwardedToggle(t *testing.T) { wantIP string }{ {name: "trust disabled records proxy address", trustForwarded: false, wantIP: "127.0.0.1"}, - {name: "trust enabled records forwarded client IP", trustForwarded: true, wantIP: "1.2.3.4"}, + {name: "legacy trust toggle cannot bypass trusted proxies", trustForwarded: true, wantIP: "127.0.0.1"}, } { t.Run(tc.name, func(t *testing.T) { cfg := &config.Config{} @@ -54,6 +53,23 @@ func TestSessionBindingContextHonorsTrustForwardedToggle(t *testing.T) { } } +func TestSessionBindingContextBoundsPersistedUserAgent(t *testing.T) { + cfg := &config.Config{} + r := gin.New() + r.Use(SessionBindingContext(cfg)) + r.GET("/t", func(c *gin.Context) { + binding := service.SessionBindingFromContext(c.Request.Context()) + require.Len(t, binding.UserAgent, maxPersistentUserAgentBytes) + require.Equal(t, binding.UserAgent, c.Request.UserAgent()) + c.Status(200) + }) + w := httptest.NewRecorder() + req := httptest.NewRequest("GET", "/t", nil) + req.Header.Set("User-Agent", strings.Repeat("u", 2048)) + r.ServeHTTP(w, req) + require.Equal(t, 200, w.Code) +} + // 未经过 SessionBindingContext 注入时(异常挂载顺序/单测直调),回退 trusted_proxies 链, // 等价于开关关闭时的历史行为。 func TestSecurityClientIPFallsBackWithoutInjectedBinding(t *testing.T) { @@ -75,8 +91,6 @@ func TestSecurityClientIPFallsBackWithoutInjectedBinding(t *testing.T) { require.Equal(t, "9.9.9.9", w.Body.String()) } -// requestSessionBinding 优先取注入值:开关开启时校验哈希必须基于注入的转发 IP 计算, -// 与 token 签发路径取值一致,否则同一客户端会被误判为指纹变化。 func TestRequestSessionBindingPrefersInjectedBinding(t *testing.T) { gin.SetMode(gin.TestMode) @@ -84,7 +98,7 @@ func TestRequestSessionBindingPrefersInjectedBinding(t *testing.T) { cfg.SetTrustForwardedIPForAPIKeyACL(true) r := gin.New() - require.NoError(t, r.SetTrustedProxies(nil)) + require.NoError(t, r.SetTrustedProxies([]string{"127.0.0.1"})) r.Use(SessionBindingContext(cfg)) r.GET("/t", func(c *gin.Context) { issued := &service.SessionBinding{IP: "1.2.3.4", UserAgent: "test-agent"} diff --git a/backend/internal/server/router.go b/backend/internal/server/router.go index 5396a94eb..47bffbef2 100644 --- a/backend/internal/server/router.go +++ b/backend/internal/server/router.go @@ -35,6 +35,7 @@ func SetupRouter( cfg *config.Config, redisClient *redis.Client, ) *gin.Engine { + middleware2.SetIngressRejectRecorder(opsService) // 缓存 iframe 页面的 origin 列表,用于动态注入 CSP frame-src var cachedFrameOrigins atomic.Pointer[[]string] emptyOrigins := []string{} @@ -55,7 +56,7 @@ func SetupRouter( // 应用中间件 r.Use(middleware2.RequestLogger()) // 将客户端 IP + UA 注入 request context,供 token 签发/会话绑定/审计日志统一读取。 - // IP 取值与 API Key IP 限制共用「信任反代传递的客户端 IP」系统开关。 + // IP 取值与 API Key IP 限制共用 server.trusted_proxies 信任链。 r.Use(middleware2.SessionBindingContext(cfg)) r.Use(middleware2.Logger()) r.Use(middleware2.CORS(cfg.CORS)) diff --git a/backend/internal/server/routes/admin.go b/backend/internal/server/routes/admin.go index 54cf0ae95..eeec9dfe3 100644 --- a/backend/internal/server/routes/admin.go +++ b/backend/internal/server/routes/admin.go @@ -235,6 +235,11 @@ func registerOpsRoutes(admin *gin.RouterGroup, h *handler.Handlers) { ops.GET("/request-errors/:id/upstream-errors", h.Admin.Ops.ListRequestErrorUpstreamErrors) ops.PUT("/request-errors/:id/resolve", h.Admin.Ops.ResolveRequestError) + // Bounded ingress-admission rejection aggregates. + ops.GET("/ingress-rejections", h.Admin.Ops.ListIngressRejects) + ops.GET("/ingress-rejections/health", h.Admin.Ops.GetIngressRejectHealth) + ops.GET("/auth-cache-invalidation/health", h.Admin.Ops.GetAuthCacheInvalidationHealth) + // Upstream errors (independent upstream failures) ops.GET("/upstream-errors", h.Admin.Ops.ListUpstreamErrors) ops.GET("/upstream-errors/:id", h.Admin.Ops.GetUpstreamError) diff --git a/backend/internal/server/routes/gateway.go b/backend/internal/server/routes/gateway.go index a1e1de260..ffac8535e 100644 --- a/backend/internal/server/routes/gateway.go +++ b/backend/internal/server/routes/gateway.go @@ -23,6 +23,7 @@ func RegisterGatewayRoutes( cfg *config.Config, ) { bodyLimit := middleware.RequestBodyLimit(cfg.Gateway.MaxBodySize) + textBodyLimit := middleware.RequestBodyLimit(cfg.Gateway.TextMaxBodySize) clientRequestID := middleware.ClientRequestID() opsErrorLogger := handler.OpsErrorLoggerMiddleware(opsService) endpointNorm := handler.InboundEndpointMiddleware() @@ -165,7 +166,7 @@ func RegisterGatewayRoutes( } h.Gateway.Responses(c) }) - gateway.POST("/alpha/search", h.OpenAIGateway.AlphaSearch) + gateway.POST("/alpha/search", textBodyLimit, h.OpenAIGateway.AlphaSearch) gateway.GET("/responses", func(c *gin.Context) { h.OpenAIGateway.ResponsesWebSocket(c) }) @@ -177,7 +178,7 @@ func RegisterGatewayRoutes( } h.Gateway.ChatCompletions(c) }) - gateway.POST("/embeddings", func(c *gin.Context) { + gateway.POST("/embeddings", textBodyLimit, func(c *gin.Context) { if getGroupPlatform(c) != service.PlatformOpenAI { service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate) c.JSON(http.StatusNotFound, gin.H{ @@ -236,7 +237,7 @@ func RegisterGatewayRoutes( } r.POST("/responses", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, responsesHandler) r.POST("/responses/*subpath", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, responsesHandler) - r.POST("/alpha/search", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, h.OpenAIGateway.AlphaSearch) + r.POST("/alpha/search", textBodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, h.OpenAIGateway.AlphaSearch) r.GET("/responses", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, func(c *gin.Context) { h.OpenAIGateway.ResponsesWebSocket(c) }) @@ -246,7 +247,7 @@ func RegisterGatewayRoutes( { codexDirect.POST("/responses", responsesHandler) codexDirect.POST("/responses/*subpath", responsesHandler) - codexDirect.POST("/alpha/search", h.OpenAIGateway.AlphaSearch) + codexDirect.POST("/alpha/search", textBodyLimit, h.OpenAIGateway.AlphaSearch) codexDirect.GET("/responses", func(c *gin.Context) { h.OpenAIGateway.ResponsesWebSocket(c) }) @@ -260,7 +261,7 @@ func RegisterGatewayRoutes( } h.Gateway.ChatCompletions(c) }) - r.POST("/embeddings", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, func(c *gin.Context) { + r.POST("/embeddings", textBodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, func(c *gin.Context) { if getGroupPlatform(c) != service.PlatformOpenAI { service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate) c.JSON(http.StatusNotFound, gin.H{ diff --git a/backend/internal/server/routes/ops_ingress_reject_routes_test.go b/backend/internal/server/routes/ops_ingress_reject_routes_test.go new file mode 100644 index 000000000..c8c5a8766 --- /dev/null +++ b/backend/internal/server/routes/ops_ingress_reject_routes_test.go @@ -0,0 +1,53 @@ +package routes + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/handler" + adminhandler "github.com/Wei-Shaw/sub2api/internal/handler/admin" + servermiddleware "github.com/Wei-Shaw/sub2api/internal/server/middleware" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestIngressRejectAdminRoutesRequireAdminAuthentication(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + handlers := &handler.Handlers{Admin: &handler.AdminHandlers{Ops: adminhandler.NewOpsHandler(nil)}} + adminAuth := servermiddleware.AdminAuthMiddleware(func(c *gin.Context) { + if c.GetHeader("Authorization") == "" { + servermiddleware.AbortWithError(c, http.StatusUnauthorized, "UNAUTHORIZED", "Authorization required") + return + } + servermiddleware.AbortWithError(c, http.StatusForbidden, "FORBIDDEN", "Admin access required") + }) + auditLog := servermiddleware.AuditLogMiddleware(func(c *gin.Context) { c.Next() }) + stepUp := servermiddleware.StepUpAuthMiddleware(func(c *gin.Context) { c.Next() }) + RegisterAdminRoutes(router.Group("/api/v1"), handlers, adminAuth, auditLog, stepUp, nil) + + for _, path := range []string{ + "/api/v1/admin/ops/ingress-rejections", + "/api/v1/admin/ops/ingress-rejections/health", + } { + for _, tc := range []struct { + name string + auth string + wantStatus int + }{ + {name: "unauthenticated", wantStatus: http.StatusUnauthorized}, + {name: "non-admin", auth: "Bearer user-token", wantStatus: http.StatusForbidden}, + } { + t.Run(path+"/"+tc.name, func(t *testing.T) { + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodGet, path, nil) + if tc.auth != "" { + request.Header.Set("Authorization", tc.auth) + } + router.ServeHTTP(recorder, request) + require.Equal(t, tc.wantStatus, recorder.Code) + }) + } + } +} diff --git a/backend/internal/service/account.go b/backend/internal/service/account.go index 724925631..5b3707d2a 100644 --- a/backend/internal/service/account.go +++ b/backend/internal/service/account.go @@ -1423,8 +1423,12 @@ func (a *Account) SupportsOpenAIEndpointCapability(capability OpenAIEndpointCapa case OpenAIEndpointCapabilityChatCompletions: return true case OpenAIEndpointCapabilityGrokMediaGeneration: - eligible, _ := a.GrokMediaGenerationEligibility() - return eligible + eligible, reason := a.GrokMediaGenerationEligibility() + // Unobserved OAuth accounts remain scheduler candidates only so the + // request path can run the billing probe before forwarding. The + // forwarding gate itself fails closed if that probe is unavailable or + // cannot produce positive paid-entitlement evidence. + return eligible || reason == "billing_unobserved" default: return false } @@ -1469,9 +1473,9 @@ func (a *Account) SupportsOpenAIEndpointCapability(capability OpenAIEndpointCapa } // GrokMediaGenerationEligibility reports whether a Grok account may receive -// new image/video generation requests. Missing observations preserve legacy -// routing; operators can fail closed for known-bad accounts with the explicit -// override. A successful override takes precedence over stale probe data. +// new image/video generation requests. OAuth media fails closed unless billing +// observations provide positive paid-entitlement evidence. An explicit +// operator override takes precedence over probe data. func (a *Account) GrokMediaGenerationEligibility() (bool, string) { if a == nil || !a.IsGrok() { return false, "not_grok" @@ -1488,11 +1492,17 @@ func (a *Account) GrokMediaGenerationEligibility() (bool, string) { billing, err := grokBillingSnapshotFromExtra(a.Extra) if err != nil || billing == nil { - return true, "billing_unobserved" + return false, "billing_unobserved" } if billing.StatusCode == 403 || billing.WeeklyStatusCode == 403 || billing.MonthlyStatusCode == 403 { return false, "billing_forbidden" } + if isKnownGrokFreeAccount(a) { + return false, "billing_free_tier" + } + if !grokBillingHasAuthoritativeQuota(billing) { + return false, "billing_inconclusive" + } return true, "eligible" } diff --git a/backend/internal/service/account_grok_media_eligibility_test.go b/backend/internal/service/account_grok_media_eligibility_test.go index be7b4759a..a79f9c02e 100644 --- a/backend/internal/service/account_grok_media_eligibility_test.go +++ b/backend/internal/service/account_grok_media_eligibility_test.go @@ -13,6 +13,7 @@ import ( ) func TestGrokMediaGenerationEligibility(t *testing.T) { + weeklyUsagePercent := 12.5 forbiddenBilling := &xai.BillingSummary{ StatusCode: http.StatusForbidden, WeeklyStatusCode: http.StatusForbidden, @@ -20,9 +21,24 @@ func TestGrokMediaGenerationEligibility(t *testing.T) { } weeklyAllowance := &xai.BillingSummary{ PeriodType: "weekly", + UsagePercent: &weeklyUsagePercent, StatusCode: http.StatusOK, WeeklyStatusCode: http.StatusOK, } + freeBilling := &xai.BillingSummary{ + PeriodType: "monthly", + 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"}, + } weeklyForbidden := &xai.BillingSummary{ StatusCode: http.StatusOK, WeeklyStatusCode: http.StatusForbidden, @@ -43,12 +59,14 @@ func TestGrokMediaGenerationEligibility(t *testing.T) { {name: "nil account", account: nil, want: false, wantReason: "not_grok"}, {name: "non grok account", account: &Account{Platform: PlatformOpenAI}, want: false, wantReason: "not_grok"}, {name: "non oauth grok account stays eligible", account: &Account{Platform: PlatformGrok, Type: AccountTypeAPIKey}, want: true, wantReason: "non_oauth"}, - {name: "unobserved oauth preserves legacy routing", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth}, want: true, wantReason: "billing_unobserved"}, - {name: "weekly allowance is not treated as weekly subscription", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: weeklyAllowance}}, want: true, wantReason: "eligible"}, + {name: "unobserved oauth fails closed", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth}, want: false, wantReason: "billing_unobserved"}, + {name: "weekly paid usage is eligible without inferring from period type", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: weeklyAllowance}}, want: true, wantReason: "eligible"}, + {name: "observed free account is rejected", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: freeBilling}}, want: false, wantReason: "billing_free_tier"}, + {name: "inconclusive billing fails closed", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: inconclusiveBilling}}, want: false, wantReason: "billing_inconclusive"}, {name: "billing forbidden is rejected", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: forbiddenBilling}}, want: false, wantReason: "billing_forbidden"}, {name: "weekly billing forbidden is rejected after partial success", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: weeklyForbidden}}, want: false, wantReason: "billing_forbidden"}, {name: "monthly billing forbidden is rejected after partial success", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: monthlyForbidden}}, want: false, wantReason: "billing_forbidden"}, - {name: "malformed billing observation preserves legacy routing", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: make(chan int)}}, want: true, wantReason: "billing_unobserved"}, + {name: "malformed billing observation fails closed", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: make(chan int)}}, want: false, wantReason: "billing_unobserved"}, {name: "malformed override falls back to observations", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{GrokMediaEligibleExtraKey: "false", grokBillingExtraKey: weeklyAllowance}}, want: true, wantReason: "eligible"}, {name: "explicit disable wins", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{GrokMediaEligibleExtraKey: false}}, want: false, wantReason: "override_disabled"}, {name: "explicit enable wins over forbidden probe", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{GrokMediaEligibleExtraKey: true, grokBillingExtraKey: forbiddenBilling}}, want: true, wantReason: "override_enabled"}, @@ -63,6 +81,24 @@ func TestGrokMediaGenerationEligibility(t *testing.T) { } } +func TestGrokMediaCapabilityKeepsOnlyUnobservedOAuthAsProbeCandidate(t *testing.T) { + unobserved := &Account{Platform: PlatformGrok, Type: AccountTypeOAuth} + eligible, reason := unobserved.GrokMediaGenerationEligibility() + require.False(t, eligible) + require.Equal(t, "billing_unobserved", reason) + 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)) +} + func TestGrokMediaCapabilityFiltersOnlyGeneration(t *testing.T) { account := &Account{ ID: 1, diff --git a/backend/internal/service/account_service.go b/backend/internal/service/account_service.go index b7fe0a68c..c32821000 100644 --- a/backend/internal/service/account_service.go +++ b/backend/internal/service/account_service.go @@ -91,6 +91,13 @@ type AccountRepository interface { ListSchedulableByGroupIDAndPlatforms(ctx context.Context, groupID int64, platforms []string) ([]Account, error) ListSchedulableUngroupedByPlatform(ctx context.Context, platform string) ([]Account, error) ListSchedulableUngroupedByPlatforms(ctx context.Context, platforms []string) ([]Account, error) + // ListModelAvailabilityCandidates returns accounts that are enabled by + // persistent configuration (active + schedulable) for model-support + // diagnosis. It deliberately does not filter transient runtime state such + // as rate-limit, overload, temporary-unschedulable, or expiry windows. + // When groupID is nil, includeGrouped controls whether the query scans all + // matching accounts or only accounts without a group binding. + ListModelAvailabilityCandidates(ctx context.Context, groupID *int64, platforms []string, includeGrouped bool) ([]Account, error) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error SetModelRateLimit(ctx context.Context, id int64, scope string, resetAt time.Time, reason ...string) error @@ -144,6 +151,7 @@ type AccountBulkUpdate struct { Schedulable *bool Credentials map[string]any Extra map[string]any + ProbeEnabled *bool } // CreateAccountRequest 创建账号请求 diff --git a/backend/internal/service/account_service_delete_test.go b/backend/internal/service/account_service_delete_test.go index ee6163239..8d93e6a93 100644 --- a/backend/internal/service/account_service_delete_test.go +++ b/backend/internal/service/account_service_delete_test.go @@ -159,6 +159,10 @@ func (s *accountRepoStub) ListSchedulableUngroupedByPlatforms(ctx context.Contex panic("unexpected ListSchedulableUngroupedByPlatforms call") } +func (s *accountRepoStub) ListModelAvailabilityCandidates(ctx context.Context, groupID *int64, platforms []string, includeGrouped bool) ([]Account, error) { + panic("unexpected ListModelAvailabilityCandidates call") +} + func (s *accountRepoStub) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error { panic("unexpected SetRateLimited call") } diff --git a/backend/internal/service/admin_account.go b/backend/internal/service/admin_account.go index 899679221..ae3a9da36 100644 --- a/backend/internal/service/admin_account.go +++ b/backend/internal/service/admin_account.go @@ -469,6 +469,15 @@ func buildAccountForCreate(input *CreateAccountInput, accountExtra map[string]an Status: StatusActive, Schedulable: true, } + if input.ProbeEnabled != nil && *input.ProbeEnabled { + if !isUpstreamBillingProbeAccount(account) { + return nil, ErrUpstreamBillingProbeAccountInvalid + } + if account.Extra == nil { + account.Extra = make(map[string]any) + } + account.Extra[UpstreamBillingProbeEnabledExtraKey] = true + } // 预计算固定时间重置的下次重置时间 if account.Extra != nil { if err := ValidateQuotaResetConfig(account.Extra); err != nil { @@ -842,7 +851,7 @@ func (s *adminServiceImpl) UpdateAccountExtra(ctx context.Context, id int64, upd // BulkUpdateAccounts updates multiple accounts in one request. // It merges credentials/extra keys instead of overwriting the whole object. func (s *adminServiceImpl) BulkUpdateAccounts(ctx context.Context, input *BulkUpdateAccountsInput) (*BulkUpdateAccountsResult, error) { - // Probe state is updated only through its dedicated endpoints. + // Managed probe state may only enter through the dedicated typed field below. delete(input.Extra, UpstreamBillingProbeEnabledExtraKey) delete(input.Extra, UpstreamBillingProbeExtraKey) @@ -874,13 +883,30 @@ func (s *adminServiceImpl) BulkUpdateAccounts(ctx context.Context, input *BulkUp // 预取所有目标账号,供凭据守卫/代理守卫/混合渠道检查共用,避免多次 DB 查询。 var cachedTargets []*Account - if len(input.Credentials) > 0 || input.ProxyID != nil || needMixedChannelCheck || hasLongContextBillingUpdate { + if len(input.Credentials) > 0 || input.ProxyID != nil || needMixedChannelCheck || hasLongContextBillingUpdate || input.ProbeEnabled != nil { loaded, err := s.accountRepo.GetByIDs(ctx, input.AccountIDs) if err != nil { return nil, err } cachedTargets = loaded } + if input.ProbeEnabled != nil { + targetsByID := make(map[int64]*Account, len(cachedTargets)) + for _, account := range cachedTargets { + if account != nil { + targetsByID[account.ID] = account + } + } + for _, accountID := range input.AccountIDs { + account, ok := targetsByID[accountID] + if !ok { + return nil, ErrAccountNotFound + } + if !isUpstreamBillingProbeAccount(account) { + return nil, ErrUpstreamBillingProbeAccountInvalid + } + } + } if hasLongContextBillingUpdate { for _, account := range cachedTargets { if account == nil || account.Platform != PlatformOpenAI { @@ -952,8 +978,15 @@ func (s *adminServiceImpl) BulkUpdateAccounts(ctx context.Context, input *BulkUp // Prepare bulk updates for columns and JSONB fields. repoUpdates := AccountBulkUpdate{ - Credentials: input.Credentials, - Extra: input.Extra, + Credentials: input.Credentials, + Extra: input.Extra, + ProbeEnabled: input.ProbeEnabled, + } + if input.ProbeEnabled != nil { + if repoUpdates.Extra == nil { + repoUpdates.Extra = make(map[string]any) + } + repoUpdates.Extra[UpstreamBillingProbeEnabledExtraKey] = *input.ProbeEnabled } if updatesUpstreamBillingProbeIdentity(input.Credentials) || input.ProxyID != nil { if repoUpdates.Extra == nil { diff --git a/backend/internal/service/admin_account_upstream_billing_probe_test.go b/backend/internal/service/admin_account_upstream_billing_probe_test.go index 727a66cf1..88db65386 100644 --- a/backend/internal/service/admin_account_upstream_billing_probe_test.go +++ b/backend/internal/service/admin_account_upstream_billing_probe_test.go @@ -38,6 +38,32 @@ func TestCreateAccountDropsManagedUpstreamBillingProbeState(t *testing.T) { require.NotContains(t, created.Extra, UpstreamBillingProbeExtraKey) } +func TestCreateAccountAcceptsDedicatedUpstreamBillingProbeSetting(t *testing.T) { + enabled := true + repo := &upstreamBillingProbeAccountRepo{} + created, err := (&adminServiceImpl{accountRepo: repo}).CreateAccount(context.Background(), &CreateAccountInput{ + Name: "upstream", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Credentials: map[string]any{"api_key": "sk-test"}, + ProbeEnabled: &enabled, + SkipDefaultGroupBind: true, + }) + + require.NoError(t, err) + require.Equal(t, true, created.Extra[UpstreamBillingProbeEnabledExtraKey]) + + _, err = (&adminServiceImpl{accountRepo: repo}).CreateAccount(context.Background(), &CreateAccountInput{ + Name: "oauth", + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Credentials: map[string]any{"access_token": "token"}, + ProbeEnabled: &enabled, + SkipDefaultGroupBind: true, + }) + require.ErrorIs(t, err, ErrUpstreamBillingProbeAccountInvalid) +} + func TestUpdateAccountPreservesManagedUpstreamBillingProbeStateForUnrelatedEdit(t *testing.T) { accountID := int64(110) repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{ @@ -360,6 +386,63 @@ func TestBulkUpdateAccountsDropsManagedUpstreamBillingProbeState(t *testing.T) { require.NotContains(t, repo.bulkUpdates[0].Extra, UpstreamBillingProbeExtraKey) } +func TestBulkUpdateAccountsAcceptsDedicatedUpstreamBillingProbeSetting(t *testing.T) { + for _, enabled := range []bool{true, false} { + t.Run(map[bool]string{true: "enable", false: "disable"}[enabled], func(t *testing.T) { + repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{ + 1: {ID: 1, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}, + 2: {ID: 2, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}, + }} + + result, err := (&adminServiceImpl{accountRepo: repo}).BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{ + AccountIDs: []int64{1, 2}, + ProbeEnabled: &enabled, + }) + + require.NoError(t, err) + require.Equal(t, 2, result.Success) + require.Len(t, repo.bulkUpdates, 1) + require.Equal(t, enabled, repo.bulkUpdates[0].Extra[UpstreamBillingProbeEnabledExtraKey]) + require.NotNil(t, repo.bulkUpdates[0].ProbeEnabled) + require.Equal(t, enabled, *repo.bulkUpdates[0].ProbeEnabled) + }) + } +} + +func TestBulkUpdateAccountsRejectsProbeSettingForIneligibleTargetBeforeWrite(t *testing.T) { + for _, enabled := range []bool{true, false} { + t.Run(map[bool]string{true: "enable", false: "disable"}[enabled], func(t *testing.T) { + repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{ + 1: {ID: 1, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}, + 2: {ID: 2, Platform: PlatformOpenAI, Type: AccountTypeOAuth}, + }} + + _, err := (&adminServiceImpl{accountRepo: repo}).BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{ + AccountIDs: []int64{1, 2}, + ProbeEnabled: &enabled, + }) + + require.ErrorIs(t, err, ErrUpstreamBillingProbeAccountInvalid) + require.Empty(t, repo.bulkUpdates) + }) + } +} + +func TestBulkUpdateAccountsRejectsProbeSettingWhenTargetIsMissing(t *testing.T) { + enabled := true + repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{ + 1: {ID: 1, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}, + }} + + _, err := (&adminServiceImpl{accountRepo: repo}).BulkUpdateAccounts(context.Background(), &BulkUpdateAccountsInput{ + AccountIDs: []int64{1, 2}, + ProbeEnabled: &enabled, + }) + + require.ErrorIs(t, err, ErrAccountNotFound) + require.Empty(t, repo.bulkUpdates) +} + func TestBulkUpdateAccountsInvalidatesProbeSnapshotForIdentityCredentials(t *testing.T) { repo := &upstreamBillingProbeAccountRepo{} input := &BulkUpdateAccountsInput{ diff --git a/backend/internal/service/admin_service.go b/backend/internal/service/admin_service.go index c167b6be8..013cb0404 100644 --- a/backend/internal/service/admin_service.go +++ b/backend/internal/service/admin_service.go @@ -329,6 +329,7 @@ type CreateAccountInput struct { GroupIDs []int64 ExpiresAt *int64 AutoPauseOnExpired *bool + ProbeEnabled *bool // SkipDefaultGroupBind prevents auto-binding to platform default group when GroupIDs is empty. SkipDefaultGroupBind bool // SkipMixedChannelCheck skips the mixed channel risk check when binding groups. @@ -378,6 +379,7 @@ type BulkUpdateAccountsInput struct { GroupIDs *[]int64 Credentials map[string]any Extra map[string]any + ProbeEnabled *bool // SkipMixedChannelCheck skips the mixed channel risk check when binding groups. // This should only be set when the caller has explicitly confirmed the risk. SkipMixedChannelCheck bool diff --git a/backend/internal/service/antigravity_gateway_service.go b/backend/internal/service/antigravity_gateway_service.go index 9e0cc804d..b835a8410 100644 --- a/backend/internal/service/antigravity_gateway_service.go +++ b/backend/internal/service/antigravity_gateway_service.go @@ -198,17 +198,18 @@ func (s *AntigravityGatewayService) getUpstreamErrorDetail(body []byte) string { } // checkErrorPolicy nil 安全的包装 -func (s *AntigravityGatewayService) checkErrorPolicy(ctx context.Context, account *Account, statusCode int, body []byte) ErrorPolicyResult { +func (s *AntigravityGatewayService) checkErrorPolicy(ctx context.Context, account *Account, statusCode int, body []byte, requestedModel ...string) ErrorPolicyResult { if s.rateLimitService == nil { return ErrorPolicyNone } - return s.rateLimitService.CheckErrorPolicy(ctx, account, statusCode, body) + return s.rateLimitService.CheckErrorPolicy(ctx, account, statusCode, body, firstRequestedModel(requestedModel)) } // applyErrorPolicy 应用错误策略结果,返回是否应终止当前循环及应返回的状态码。 // ErrorPolicySkipped 时 outStatus 为 500(前端约定:未命中的错误返回 500)。 func (s *AntigravityGatewayService) applyErrorPolicy(p antigravityRetryLoopParams, statusCode int, headers http.Header, respBody []byte) (handled bool, outStatus int, retErr error) { - switch s.checkErrorPolicy(p.ctx, p.account, statusCode, respBody) { + modelKey := resolveFinalAntigravityModelKey(p.ctx, p.account, p.requestedModel) + switch s.checkErrorPolicy(p.ctx, p.account, statusCode, respBody, modelKey) { case ErrorPolicySkipped: if s.handleAntigravityModelRateLimitBeforePolicy(p, statusCode, headers, respBody) { return true, statusCode, nil diff --git a/backend/internal/service/api_key_auth_cache_impl.go b/backend/internal/service/api_key_auth_cache_impl.go index 5c4540865..4cf47c8de 100644 --- a/backend/internal/service/api_key_auth_cache_impl.go +++ b/backend/internal/service/api_key_auth_cache_impl.go @@ -76,32 +76,119 @@ func (c apiKeyAuthCacheConfig) jitterTTL(ttl time.Duration) time.Duration { func (s *APIKeyService) initAuthCache(cfg *config.Config) { s.authCfg = newAPIKeyAuthCacheConfig(cfg) - if !s.authCfg.l1Enabled() { - return + if s.authCfg.negativeEnabled() { + negativeSize := defaultNegativeAuthCacheSize + if s.authCfg.l1Size > 0 && s.authCfg.l1Size < negativeSize { + negativeSize = s.authCfg.l1Size + } + cache, err := ristretto.NewCache(&ristretto.Config{ + NumCounters: int64(negativeSize) * 10, + MaxCost: int64(negativeSize), + BufferItems: 64, + }) + if err == nil { + s.authNegativeCacheL1 = cache + } } - cache, err := ristretto.NewCache(&ristretto.Config{ - NumCounters: int64(s.authCfg.l1Size) * 10, - MaxCost: int64(s.authCfg.l1Size), - BufferItems: 64, - }) - if err != nil { - return + if s.authCfg.l1Enabled() { + cache, err := ristretto.NewCache(&ristretto.Config{ + NumCounters: int64(s.authCfg.l1Size) * 10, + MaxCost: int64(s.authCfg.l1Size), + BufferItems: 64, + }) + if err == nil { + s.authCacheL1 = cache + } } - s.authCacheL1 = cache } // StartAuthCacheInvalidationSubscriber starts the Pub/Sub subscriber for L1 cache invalidation. // This should be called after the service is fully initialized. func (s *APIKeyService) StartAuthCacheInvalidationSubscriber(ctx context.Context) { - if s.cache == nil || s.authCacheL1 == nil { + if s.cache == nil || (s.authCacheL1 == nil && s.authNegativeCacheL1 == nil) { return } - if err := s.cache.SubscribeAuthCacheInvalidation(ctx, func(cacheKey string) { - s.authCacheL1.Del(cacheKey) - }); err != nil { - // Log but don't fail - L1 cache will still work, just without cross-instance invalidation - slog.Warn("failed to start auth cache invalidation subscriber", "error", err) + s.authInvalidationStart.Do(func() { + subscriberCtx, cancel := context.WithCancel(ctx) + subscriberCtx = withAuthCacheSubscriptionReady(subscriberCtx, func() { + s.authInvalidationConnected.Store(true) + }) + s.authInvalidationCancel = cancel + s.authInvalidationWG.Add(1) + go func() { + defer s.authInvalidationWG.Done() + backoff := time.Second + for { + err := s.cache.SubscribeAuthCacheInvalidation(subscriberCtx, func(cacheKey string) { + s.invalidateLocalAuthCache(cacheKey) + }) + wasConnected := s.authInvalidationConnected.Swap(false) + if subscriberCtx.Err() != nil { + return + } + if wasConnected { + backoff = time.Second + } + s.authInvalidationFailures.Add(1) + if err == nil { + err = errors.New("auth cache invalidation subscription closed") + } + slog.Warn("failed to start auth cache invalidation subscriber; retrying", "error", err, "retry_in", backoff) + timer := time.NewTimer(backoff) + select { + case <-subscriberCtx.Done(): + timer.Stop() + return + case <-timer.C: + } + if backoff < 30*time.Second { + backoff *= 2 + if backoff > 30*time.Second { + backoff = 30 * time.Second + } + } + } + }() + }) +} + +func (s *APIKeyService) invalidateLocalAuthCache(cacheKey string) { + if s == nil { + return } + if s.authCacheL1 != nil { + s.authCacheL1.Del(cacheKey) + } + if s.authNegativeCacheL1 != nil { + s.authNegativeCacheL1.Del(cacheKey) + } +} + +type AuthCacheInvalidationSubscriberHealth struct { + Connected bool `json:"connected"` + Failures uint64 `json:"failures"` +} + +func (s *APIKeyService) AuthCacheInvalidationSubscriberHealth() AuthCacheInvalidationSubscriberHealth { + if s == nil { + return AuthCacheInvalidationSubscriberHealth{} + } + return AuthCacheInvalidationSubscriberHealth{ + Connected: s.authInvalidationConnected.Load(), + Failures: s.authInvalidationFailures.Load(), + } +} + +func (s *APIKeyService) StopAuthCacheInvalidationSubscriber() { + if s == nil { + return + } + s.authInvalidationStop.Do(func() { + if s.authInvalidationCancel != nil { + s.authInvalidationCancel() + } + s.authInvalidationWG.Wait() + }) } func (s *APIKeyService) authCacheKey(key string) string { @@ -117,6 +204,13 @@ func (s *APIKeyService) getAuthCacheEntry(ctx context.Context, cacheKey string) } } } + if s.authNegativeCacheL1 != nil { + if val, ok := s.authNegativeCacheL1.Get(cacheKey); ok { + if entry, ok := val.(*APIKeyAuthCacheEntry); ok && entry.NotFound { + return entry, true + } + } + } if s.cache == nil || !s.authCfg.l2Enabled() { return nil, false } @@ -129,13 +223,19 @@ func (s *APIKeyService) getAuthCacheEntry(ctx context.Context, cacheKey string) } func (s *APIKeyService) setAuthCacheL1(cacheKey string, entry *APIKeyAuthCacheEntry) { - if s.authCacheL1 == nil || entry == nil { + if entry == nil { + return + } + if entry.NotFound { + if s.authNegativeCacheL1 != nil && s.authCfg.negativeTTL > 0 { + _ = s.authNegativeCacheL1.SetWithTTL(cacheKey, entry, 1, s.authCfg.jitterTTL(s.authCfg.negativeTTL)) + } + return + } + if s.authCacheL1 == nil { return } ttl := s.authCfg.l1TTL - if entry.NotFound && s.authCfg.negativeTTL > 0 && s.authCfg.negativeTTL < ttl { - ttl = s.authCfg.negativeTTL - } ttl = s.authCfg.jitterTTL(ttl) _ = s.authCacheL1.SetWithTTL(cacheKey, entry, 1, ttl) } @@ -155,6 +255,9 @@ func (s *APIKeyService) deleteAuthCache(ctx context.Context, cacheKey string) { if s.authCacheL1 != nil { s.authCacheL1.Del(cacheKey) } + if s.authNegativeCacheL1 != nil { + s.authNegativeCacheL1.Del(cacheKey) + } if s.cache == nil { return } @@ -164,12 +267,15 @@ func (s *APIKeyService) deleteAuthCache(ctx context.Context, cacheKey string) { } func (s *APIKeyService) loadAuthCacheEntry(ctx context.Context, key, cacheKey string) (*APIKeyAuthCacheEntry, error) { - apiKey, err := s.apiKeyRepo.GetByKeyForAuth(ctx, key) + apiKey, err := s.lookupAPIKeyForAuth(ctx, key) if err != nil { if errors.Is(err, ErrAPIKeyNotFound) { entry := &APIKeyAuthCacheEntry{NotFound: true} if s.authCfg.negativeEnabled() { - s.setAuthCacheEntry(ctx, cacheKey, entry, s.authCfg.negativeTTL) + // Invalid keys are attacker-controlled and high-cardinality. Keep their + // negative entries in the bounded process-local cache; do not amplify + // random-key scans into Redis writes on every instance. + s.setAuthCacheL1(cacheKey, entry) } return entry, nil } @@ -185,6 +291,30 @@ func (s *APIKeyService) loadAuthCacheEntry(ctx context.Context, key, cacheKey st return entry, nil } +func (s *APIKeyService) lookupAPIKeyForAuth(ctx context.Context, key string) (*APIKey, error) { + if s == nil || s.apiKeyRepo == nil { + return nil, ErrAPIKeyNotFound + } + if s.authLookupSlots == nil { + return s.apiKeyRepo.GetByKeyForAuth(ctx, key) + } + s.authLookupTotal.Add(1) + select { + case s.authLookupSlots <- struct{}{}: + s.authLookupInFlight.Add(1) + defer func() { + s.authLookupInFlight.Add(-1) + <-s.authLookupSlots + }() + case <-ctx.Done(): + return nil, ctx.Err() + default: + s.authLookupRejected.Add(1) + return nil, ErrAPIKeyAuthOverloaded + } + return s.apiKeyRepo.GetByKeyForAuth(ctx, key) +} + func (s *APIKeyService) applyAuthCacheEntry(key string, entry *APIKeyAuthCacheEntry) (*APIKey, bool, error) { if entry == nil { return nil, false, nil diff --git a/backend/internal/service/api_key_service.go b/backend/internal/service/api_key_service.go index b13ca6b3d..3b15d3174 100644 --- a/backend/internal/service/api_key_service.go +++ b/backend/internal/service/api_key_service.go @@ -10,6 +10,7 @@ import ( "strconv" "strings" "sync" + "sync/atomic" "time" "github.com/Wei-Shaw/sub2api/internal/config" @@ -22,13 +23,14 @@ import ( ) var ( - ErrAPIKeyNotFound = infraerrors.NotFound("API_KEY_NOT_FOUND", "api key not found") - 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") - 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") + ErrAPIKeyNotFound = infraerrors.NotFound("API_KEY_NOT_FOUND", "api key not found") + 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") + 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") + ErrAPIKeyAuthOverloaded = infraerrors.ServiceUnavailable("API_KEY_AUTH_OVERLOADED", "api key authentication is temporarily overloaded") + ErrInvalidIPPattern = infraerrors.BadRequest("INVALID_IP_PATTERN", "invalid IP or CIDR pattern") // ErrAPIKeyExpired = infraerrors.Forbidden("API_KEY_EXPIRED", "api key has expired") ErrAPIKeyExpired = infraerrors.Forbidden("API_KEY_EXPIRED", "api key 已过期") // ErrAPIKeyQuotaExhausted = infraerrors.TooManyRequests("API_KEY_QUOTA_EXHAUSTED", "api key quota exhausted") @@ -41,6 +43,9 @@ var ( ) const ( + MaxAPIKeyCredentialBytes = 128 + defaultAuthLookupConcurrency = 64 + defaultNegativeAuthCacheSize = 16384 apiKeyMaxErrorsPerHour = 20 apiKeyLastUsedMinTouch = 30 * time.Second apiKeySortCurrentConcurrency = "current_concurrency" @@ -58,7 +63,9 @@ type APIKeyRepository interface { GetByKeyForAuth(ctx context.Context, key string) (*APIKey, error) Update(ctx context.Context, key *APIKey) error Delete(ctx context.Context, id int64) error - // DeleteWithAudit 在同一事务内先写 deleted_api_key_audits 审计、再软删除该 key。 + // DeleteWithAudit keeps the legacy interface name for rolling-upgrade compatibility. + // Implementations must tombstone the key and soft-delete it atomically without + // retaining the deleted credential material. DeleteWithAudit(ctx context.Context, id int64) error ListByUserID(ctx context.Context, userID int64, params pagination.PaginationParams, filters APIKeyListFilters) ([]APIKey, *pagination.PaginationResult, error) @@ -149,6 +156,20 @@ type APIKeyCache interface { SubscribeAuthCacheInvalidation(ctx context.Context, handler func(cacheKey string)) error } +type authCacheSubscriptionReadyKey struct{} + +func withAuthCacheSubscriptionReady(ctx context.Context, ready func()) context.Context { + return context.WithValue(ctx, authCacheSubscriptionReadyKey{}, ready) +} + +// NotifyAuthCacheSubscriptionReady lets cache implementations report that the +// server acknowledged the subscription without widening the public cache API. +func NotifyAuthCacheSubscriptionReady(ctx context.Context) { + if ready, ok := ctx.Value(authCacheSubscriptionReadyKey{}).(func()); ok && ready != nil { + ready() + } +} + // APIKeyAuthCacheInvalidator 提供认证缓存失效能力 type APIKeyAuthCacheInvalidator interface { InvalidateAuthCacheByKey(ctx context.Context, key string) @@ -202,20 +223,51 @@ type RateLimitCacheInvalidator interface { } type APIKeyService struct { - apiKeyRepo APIKeyRepository - userRepo UserRepository - groupRepo GroupRepository - userSubRepo UserSubscriptionRepository - userGroupRateRepo UserGroupRateRepository - cache APIKeyCache - rateLimitCacheInvalid RateLimitCacheInvalidator // optional: invalidate Redis rate limit cache - concurrencyService *ConcurrencyService - cfg *config.Config - authCacheL1 *ristretto.Cache - authCfg apiKeyAuthCacheConfig - authGroup singleflight.Group - lastUsedTouchL1 sync.Map // keyID -> nextAllowedAt(time.Time) - lastUsedTouchSF singleflight.Group + apiKeyRepo APIKeyRepository + userRepo UserRepository + groupRepo GroupRepository + userSubRepo UserSubscriptionRepository + userGroupRateRepo UserGroupRateRepository + cache APIKeyCache + rateLimitCacheInvalid RateLimitCacheInvalidator // optional: invalidate Redis rate limit cache + concurrencyService *ConcurrencyService + cfg *config.Config + authCacheL1 *ristretto.Cache + authNegativeCacheL1 *ristretto.Cache + authCfg apiKeyAuthCacheConfig + authGroup singleflight.Group + authLookupSlots chan struct{} + authLookupTotal atomic.Uint64 + authLookupRejected atomic.Uint64 + authLookupInFlight atomic.Int64 + invalidAuthAbuse *invalidAuthAbuseLimiter + authInvalidationStart sync.Once + authInvalidationStop sync.Once + authInvalidationCancel context.CancelFunc + authInvalidationWG sync.WaitGroup + authInvalidationConnected atomic.Bool + authInvalidationFailures atomic.Uint64 + lastUsedTouchL1 sync.Map // keyID -> nextAllowedAt(time.Time) + lastUsedTouchSF singleflight.Group +} + +type APIKeyAuthLookupMetrics struct { + Total uint64 `json:"total"` + Rejected uint64 `json:"rejected"` + InFlight int64 `json:"in_flight"` + Capacity int `json:"capacity"` +} + +func (s *APIKeyService) AuthLookupMetrics() APIKeyAuthLookupMetrics { + if s == nil { + return APIKeyAuthLookupMetrics{} + } + return APIKeyAuthLookupMetrics{ + Total: s.authLookupTotal.Load(), + Rejected: s.authLookupRejected.Load(), + InFlight: s.authLookupInFlight.Load(), + Capacity: cap(s.authLookupSlots), + } } // NewAPIKeyService 创建API Key服务实例 @@ -238,6 +290,12 @@ func NewAPIKeyService( cfg: cfg, } svc.initAuthCache(cfg) + lookupConcurrency := defaultAuthLookupConcurrency + if cfg != nil && cfg.APIKeyAuth.LookupConcurrency > 0 { + lookupConcurrency = cfg.APIKeyAuth.LookupConcurrency + } + svc.authLookupSlots = make(chan struct{}, lookupConcurrency) + svc.invalidAuthAbuse = newInvalidAuthAbuseLimiter(cfg) return svc } @@ -581,6 +639,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 len(key) == 0 || len(key) > MaxAPIKeyCredentialBytes { + return nil, ErrAPIKeyNotFound + } cacheKey := s.authCacheKey(key) if entry, ok := s.getAuthCacheEntry(ctx, cacheKey); ok { @@ -622,7 +683,7 @@ func (s *APIKeyService) GetByKey(ctx context.Context, key string) (*APIKey, erro } } - apiKey, err := s.apiKeyRepo.GetByKeyForAuth(ctx, key) + apiKey, err := s.lookupAPIKeyForAuth(ctx, key) if err != nil { return nil, fmt.Errorf("get api key: %w", err) } diff --git a/backend/internal/service/api_key_service_cache_test.go b/backend/internal/service/api_key_service_cache_test.go index a1dfbcb01..6dfa05f8b 100644 --- a/backend/internal/service/api_key_service_cache_test.go +++ b/backend/internal/service/api_key_service_cache_test.go @@ -5,6 +5,7 @@ package service import ( "context" "errors" + "strings" "sync" "sync/atomic" "testing" @@ -509,14 +510,18 @@ func TestAPIKeyService_InvalidateAuthCacheByKey(t *testing.T) { } func TestAPIKeyService_GetByKey_CachesNegativeOnRepoMiss(t *testing.T) { + var repoCalls atomic.Int32 cache := &authCacheStub{} repo := &authRepoStub{ getByKeyForAuth: func(ctx context.Context, key string) (*APIKey, error) { + repoCalls.Add(1) return nil, ErrAPIKeyNotFound }, } cfg := &config.Config{ APIKeyAuth: config.APIKeyAuthCacheConfig{ + L1Size: 100, + L1TTLSeconds: 60, L2TTLSeconds: 60, NegativeTTLSeconds: 30, }, @@ -528,7 +533,73 @@ func TestAPIKeyService_GetByKey_CachesNegativeOnRepoMiss(t *testing.T) { _, err := svc.GetByKey(context.Background(), "missing") require.ErrorIs(t, err, ErrAPIKeyNotFound) - require.Len(t, cache.setAuthKeys, 1) + require.Empty(t, cache.setAuthKeys, "attacker-controlled misses must not be written to Redis") + svc.authNegativeCacheL1.Wait() + _, err = svc.GetByKey(context.Background(), "missing") + require.ErrorIs(t, err, ErrAPIKeyNotFound) + require.Equal(t, int32(1), repoCalls.Load()) +} + +func TestAPIKeyService_GetByKeyRejectsInvalidLengthBeforeCaches(t *testing.T) { + var cacheCalls atomic.Int32 + cache := &authCacheStub{getAuthCache: func(context.Context, string) (*APIKeyAuthCacheEntry, error) { + cacheCalls.Add(1) + return nil, redis.Nil + }} + repo := &authRepoStub{getByKeyForAuth: func(context.Context, string) (*APIKey, error) { + t.Fatal("invalid credential reached repository") + return nil, nil + }} + svc := NewAPIKeyService(repo, nil, nil, nil, nil, cache, &config.Config{APIKeyAuth: config.APIKeyAuthCacheConfig{L2TTLSeconds: 60}}) + + for _, key := range []string{"", strings.Repeat("x", MaxAPIKeyCredentialBytes+1)} { + _, err := svc.GetByKey(context.Background(), key) + require.ErrorIs(t, err, ErrAPIKeyNotFound) + } + require.Zero(t, cacheCalls.Load()) +} + +func TestAPIKeyService_GetByKeyAllowsMaximumLength(t *testing.T) { + key := strings.Repeat("x", MaxAPIKeyCredentialBytes) + var repoCalls atomic.Int32 + repo := &authRepoStub{getByKeyForAuth: func(_ context.Context, got string) (*APIKey, error) { + repoCalls.Add(1) + require.Equal(t, key, got) + return nil, ErrAPIKeyNotFound + }} + svc := NewAPIKeyService(repo, nil, nil, nil, nil, nil, &config.Config{}) + _, err := svc.GetByKey(context.Background(), key) + require.ErrorIs(t, err, ErrAPIKeyNotFound) + require.Equal(t, int32(1), repoCalls.Load()) +} + +func TestAPIKeyService_AuthLookupBulkheadRejectsExcessMisses(t *testing.T) { + entered := make(chan struct{}) + release := make(chan struct{}) + repo := &authRepoStub{getByKeyForAuth: func(context.Context, string) (*APIKey, error) { + close(entered) + <-release + return nil, ErrAPIKeyNotFound + }} + svc := NewAPIKeyService(repo, nil, nil, nil, nil, nil, &config.Config{APIKeyAuth: config.APIKeyAuthCacheConfig{LookupConcurrency: 1}}) + + done := make(chan error, 1) + go func() { + _, err := svc.GetByKey(context.Background(), "first") + done <- err + }() + <-entered + + _, err := svc.GetByKey(context.Background(), "second") + require.ErrorIs(t, err, ErrAPIKeyAuthOverloaded) + metrics := svc.AuthLookupMetrics() + require.Equal(t, uint64(2), metrics.Total) + require.Equal(t, uint64(1), metrics.Rejected) + require.Equal(t, int64(1), metrics.InFlight) + require.Equal(t, 1, metrics.Capacity) + + close(release) + require.ErrorIs(t, <-done, ErrAPIKeyNotFound) } func TestAPIKeyService_GetByKey_SingleflightCollapses(t *testing.T) { diff --git a/backend/internal/service/auth_cache_invalidation_outbox.go b/backend/internal/service/auth_cache_invalidation_outbox.go new file mode 100644 index 000000000..1c2d54043 --- /dev/null +++ b/backend/internal/service/auth_cache_invalidation_outbox.go @@ -0,0 +1,293 @@ +package service + +import ( + "context" + "fmt" + "log/slog" + "math/rand/v2" + "sync" + "sync/atomic" + "time" + + "github.com/google/uuid" +) + +const ( + authInvalidationBatchSize = 100 + authInvalidationPollInterval = 500 * time.Millisecond + authInvalidationLease = 30 * time.Second + authInvalidationRedisTimeout = 2 * time.Second + authInvalidationSafetyDelay = 30 * time.Second + authInvalidationConcurrency = 16 +) + +type AuthCacheInvalidationEvent struct { + ID int64 + CacheKey string + Attempts int + Stage int + CreatedAt time.Time +} + +type AuthCacheInvalidationOutboxStats struct { + Pending int64 + OldestCreatedAt *time.Time + MaxAttempts int + LastError string +} + +type AuthCacheInvalidationOutboxRepository interface { + Claim(ctx context.Context, workerID string, limit int, lease time.Duration) ([]AuthCacheInvalidationEvent, error) + DeleteClaimed(ctx context.Context, id int64, workerID string) error + ScheduleSecondPass(ctx context.Context, id int64, workerID string, availableAt time.Time) error + RetryClaimed(ctx context.Context, id int64, workerID string, availableAt time.Time, lastError string) error + Stats(ctx context.Context) (AuthCacheInvalidationOutboxStats, error) +} + +type AuthCacheInvalidationHealth struct { + Running bool `json:"running"` + Processed uint64 `json:"processed"` + Failures uint64 `json:"failures"` + Pending int64 `json:"pending"` + OldestLag time.Duration `json:"oldest_lag"` + LastError string `json:"last_error,omitempty"` + StatsError string `json:"stats_error,omitempty"` + // HealthySLA includes the delayed safety pass. RecoverySLA is the maximum + // convergence time after Redis becomes healthy, including capped backoff. + HealthySLA time.Duration `json:"healthy_sla"` + RecoverySLA time.Duration `json:"recovery_sla"` + MaxAttempts int `json:"max_attempts"` +} + +type OpsAuthCacheInvalidationHealth struct { + Outbox AuthCacheInvalidationHealth `json:"outbox"` + Subscriber AuthCacheInvalidationSubscriberHealth `json:"subscriber"` + Lookup APIKeyAuthLookupMetrics `json:"lookup"` + InvalidAbuse InvalidAuthAbuseHealth `json:"invalid_abuse"` +} + +func (s *OpsService) GetAuthCacheInvalidationHealth(ctx context.Context) OpsAuthCacheInvalidationHealth { + if s == nil { + return OpsAuthCacheInvalidationHealth{} + } + health := OpsAuthCacheInvalidationHealth{} + if s.authCacheInvalidationWorker != nil { + health.Outbox = s.authCacheInvalidationWorker.Health(ctx) + } + if s.apiKeyService != nil { + health.Subscriber = s.apiKeyService.AuthCacheInvalidationSubscriberHealth() + health.Lookup = s.apiKeyService.AuthLookupMetrics() + health.InvalidAbuse = s.apiKeyService.InvalidAuthAbuseHealth() + } + return health +} + +type AuthCacheInvalidationWorker struct { + repo AuthCacheInvalidationOutboxRepository + cache APIKeyCache + local *APIKeyService + workerID string + ctx context.Context + cancel context.CancelFunc + wg sync.WaitGroup + start sync.Once + stop sync.Once + running atomic.Bool + processed atomic.Uint64 + failures atomic.Uint64 + lastError atomic.Value +} + +func NewAuthCacheInvalidationWorker(repo AuthCacheInvalidationOutboxRepository, cache APIKeyCache, local ...*APIKeyService) *AuthCacheInvalidationWorker { + ctx, cancel := context.WithCancel(context.Background()) + w := &AuthCacheInvalidationWorker{ + repo: repo, cache: cache, workerID: uuid.NewString(), ctx: ctx, cancel: cancel, + } + if len(local) > 0 { + w.local = local[0] + } + w.lastError.Store("") + return w +} + +func (w *AuthCacheInvalidationWorker) Start() { + if w == nil || w.repo == nil || w.cache == nil { + return + } + w.start.Do(func() { + w.running.Store(true) + w.wg.Add(1) + go w.run() + }) +} + +func (w *AuthCacheInvalidationWorker) Stop() { + if w == nil { + return + } + w.stop.Do(func() { + w.cancel() + w.wg.Wait() + w.running.Store(false) + }) +} + +func (w *AuthCacheInvalidationWorker) run() { + defer w.wg.Done() + defer w.running.Store(false) + ticker := time.NewTicker(authInvalidationPollInterval) + defer ticker.Stop() + for { + if err := w.processBatch(w.ctx); err != nil && w.ctx.Err() == nil { + w.recordFailure(err) + } + select { + case <-w.ctx.Done(): + return + case <-ticker.C: + } + } +} + +func (w *AuthCacheInvalidationWorker) processBatch(ctx context.Context) error { + events, err := w.repo.Claim(ctx, w.workerID, authInvalidationBatchSize, authInvalidationLease) + if err != nil { + return fmt.Errorf("claim auth cache invalidations: %w", err) + } + semaphore := make(chan struct{}, authInvalidationConcurrency) + var wg sync.WaitGroup + for i := range events { + select { + case <-ctx.Done(): + wg.Wait() + return ctx.Err() + case semaphore <- struct{}{}: + } + wg.Add(1) + go func(event AuthCacheInvalidationEvent) { + defer wg.Done() + defer func() { <-semaphore }() + w.processEvent(ctx, event) + }(events[i]) + } + wg.Wait() + return nil +} + +func (w *AuthCacheInvalidationWorker) processEvent(parent context.Context, event AuthCacheInvalidationEvent) { + if w.local != nil { + w.local.invalidateLocalAuthCache(event.CacheKey) + } + ctx, cancel := context.WithTimeout(parent, authInvalidationRedisTimeout) + err := w.cache.DeleteAuthCache(ctx, event.CacheKey) + if err == nil { + err = w.cache.PublishAuthCacheInvalidation(ctx, event.CacheKey) + } + cancel() + if err != nil { + w.recordFailure(err) + retryAt := time.Now().UTC().Add(authInvalidationRetryDelay(event.Attempts + 1)) + retryCtx, retryCancel := context.WithTimeout(context.Background(), 2*time.Second) + retryErr := w.repo.RetryClaimed(retryCtx, event.ID, w.workerID, retryAt, boundedAuthInvalidationError(err)) + retryCancel() + if retryErr != nil { + w.recordFailure(fmt.Errorf("release failed auth invalidation %d: %w", event.ID, retryErr)) + } + return + } + if event.Stage == 0 { + nextCtx, nextCancel := context.WithTimeout(context.Background(), 2*time.Second) + err = w.repo.ScheduleSecondPass(nextCtx, event.ID, w.workerID, time.Now().UTC().Add(authInvalidationSafetyDelay)) + nextCancel() + if err != nil { + w.recordFailure(fmt.Errorf("schedule second auth invalidation pass %d: %w", event.ID, err)) + return + } + w.processed.Add(1) + w.lastError.Store("") + return + } + + ackCtx, ackCancel := context.WithTimeout(context.Background(), 2*time.Second) + err = w.repo.DeleteClaimed(ackCtx, event.ID, w.workerID) + ackCancel() + if err != nil { + w.recordFailure(fmt.Errorf("ack auth invalidation %d: %w", event.ID, err)) + return + } + w.processed.Add(1) + w.lastError.Store("") +} + +func authInvalidationRetryDelay(attempt int) time.Duration { + if attempt < 1 { + attempt = 1 + } + if attempt > 9 { + attempt = 9 + } + base := time.Second * time.Duration(1<<(attempt-1)) + return time.Duration(float64(base) * (0.8 + rand.Float64()*0.4)) +} + +func boundedAuthInvalidationError(err error) string { + if err == nil { + return "" + } + message := err.Error() + if len(message) > 1024 { + return message[:1024] + } + return message +} + +func (w *AuthCacheInvalidationWorker) recordFailure(err error) { + if err == nil { + return + } + w.failures.Add(1) + w.lastError.Store(boundedAuthInvalidationError(err)) + slog.Warn("auth cache invalidation outbox processing failed", "error", err) +} + +func (w *AuthCacheInvalidationWorker) Health(ctx context.Context) AuthCacheInvalidationHealth { + health := AuthCacheInvalidationHealth{ + HealthySLA: authInvalidationSafetyDelay + 5*time.Second, + RecoverySLA: 6 * time.Minute, + } + if w == nil { + return health + } + health.Running = w.running.Load() + health.Processed = w.processed.Load() + health.Failures = w.failures.Load() + if value := w.lastError.Load(); value != nil { + health.LastError, _ = value.(string) + } + if w.repo == nil { + return health + } + stats, err := w.repo.Stats(ctx) + if err != nil { + health.StatsError = boundedAuthInvalidationError(err) + return health + } + health.Pending = stats.Pending + health.MaxAttempts = stats.MaxAttempts + if health.LastError == "" { + health.LastError = stats.LastError + } + if stats.OldestCreatedAt != nil { + health.OldestLag = time.Since(*stats.OldestCreatedAt) + if health.OldestLag < 0 { + health.OldestLag = 0 + } + } + return health +} + +func ProvideAuthCacheInvalidationWorker(repo AuthCacheInvalidationOutboxRepository, cache APIKeyCache, apiKeyService *APIKeyService) *AuthCacheInvalidationWorker { + worker := NewAuthCacheInvalidationWorker(repo, cache, apiKeyService) + worker.Start() + return worker +} diff --git a/backend/internal/service/auth_cache_invalidation_outbox_test.go b/backend/internal/service/auth_cache_invalidation_outbox_test.go new file mode 100644 index 000000000..2f08723a6 --- /dev/null +++ b/backend/internal/service/auth_cache_invalidation_outbox_test.go @@ -0,0 +1,281 @@ +package service + +import ( + "context" + "errors" + "sync" + "testing" + "time" + + "github.com/dgraph-io/ristretto" + "github.com/stretchr/testify/require" +) + +type authInvalidationRepoStub struct { + mu sync.Mutex + events []AuthCacheInvalidationEvent + claimLimit int + scheduled []int64 + deleted []int64 + retried []int64 + retryError string + stats AuthCacheInvalidationOutboxStats + statsErr error +} + +func (r *authInvalidationRepoStub) Claim(_ context.Context, _ string, limit int, _ time.Duration) ([]AuthCacheInvalidationEvent, error) { + r.mu.Lock() + defer r.mu.Unlock() + r.claimLimit = limit + return append([]AuthCacheInvalidationEvent(nil), r.events...), nil +} +func (r *authInvalidationRepoStub) DeleteClaimed(_ context.Context, id int64, _ string) error { + r.mu.Lock() + defer r.mu.Unlock() + r.deleted = append(r.deleted, id) + return nil +} +func (r *authInvalidationRepoStub) ScheduleSecondPass(_ context.Context, id int64, _ string, _ time.Time) error { + r.mu.Lock() + defer r.mu.Unlock() + r.scheduled = append(r.scheduled, id) + return nil +} +func (r *authInvalidationRepoStub) RetryClaimed(_ context.Context, id int64, _ string, _ time.Time, lastError string) error { + r.mu.Lock() + defer r.mu.Unlock() + r.retried = append(r.retried, id) + r.retryError = lastError + return nil +} +func (r *authInvalidationRepoStub) Stats(context.Context) (AuthCacheInvalidationOutboxStats, error) { + return r.stats, r.statsErr +} + +type authInvalidationCacheStub struct { + mu sync.Mutex + deleteFn func(context.Context, string) error + publishFn func(context.Context, string) error + subscribeFn func(context.Context, func(string)) error + deleted []string + published []string +} + +func (*authInvalidationCacheStub) GetCreateAttemptCount(context.Context, int64) (int, error) { + return 0, nil +} +func (*authInvalidationCacheStub) IncrementCreateAttemptCount(context.Context, int64) error { + return nil +} +func (*authInvalidationCacheStub) DeleteCreateAttemptCount(context.Context, int64) error { return nil } +func (*authInvalidationCacheStub) IncrementDailyUsage(context.Context, string) error { return nil } +func (*authInvalidationCacheStub) SetDailyUsageExpiry(context.Context, string, time.Duration) error { + return nil +} +func (*authInvalidationCacheStub) GetAuthCache(context.Context, string) (*APIKeyAuthCacheEntry, error) { + return nil, errors.New("miss") +} +func (*authInvalidationCacheStub) SetAuthCache(context.Context, string, *APIKeyAuthCacheEntry, time.Duration) error { + return nil +} +func (c *authInvalidationCacheStub) DeleteAuthCache(ctx context.Context, key string) error { + c.mu.Lock() + c.deleted = append(c.deleted, key) + c.mu.Unlock() + if c.deleteFn != nil { + return c.deleteFn(ctx, key) + } + return nil +} +func (c *authInvalidationCacheStub) PublishAuthCacheInvalidation(ctx context.Context, key string) error { + c.mu.Lock() + c.published = append(c.published, key) + c.mu.Unlock() + if c.publishFn != nil { + return c.publishFn(ctx, key) + } + return nil +} +func (c *authInvalidationCacheStub) SubscribeAuthCacheInvalidation(ctx context.Context, handler func(string)) error { + if c.subscribeFn != nil { + return c.subscribeFn(ctx, handler) + } + return nil +} + +func TestAuthCacheInvalidationWorker_FirstPassSchedulesSafetyPass(t *testing.T) { + repo := &authInvalidationRepoStub{} + cache := &authInvalidationCacheStub{} + worker := NewAuthCacheInvalidationWorker(repo, cache) + worker.processEvent(context.Background(), AuthCacheInvalidationEvent{ID: 7, CacheKey: "hash", Stage: 0}) + require.Equal(t, []string{"hash"}, cache.deleted) + require.Equal(t, []string{"hash"}, cache.published) + require.Equal(t, []int64{7}, repo.scheduled) + require.Empty(t, repo.deleted) +} + +func TestAuthCacheInvalidationWorker_SecondPassCleansEvent(t *testing.T) { + repo := &authInvalidationRepoStub{} + cache := &authInvalidationCacheStub{} + worker := NewAuthCacheInvalidationWorker(repo, cache) + worker.processEvent(context.Background(), AuthCacheInvalidationEvent{ID: 8, CacheKey: "hash", Stage: 1}) + require.Equal(t, []int64{8}, repo.deleted) + require.Equal(t, uint64(1), worker.Health(context.Background()).Processed) +} + +func TestAuthCacheInvalidationWorker_RetriesRedisAndPublishFailures(t *testing.T) { + for _, tc := range []struct { + name string + deleteErr error + publishErr error + published int + }{ + {name: "redis down", deleteErr: errors.New("redis unavailable")}, + {name: "publish failure after delete", publishErr: errors.New("publish failed"), published: 1}, + } { + t.Run(tc.name, func(t *testing.T) { + repo := &authInvalidationRepoStub{} + cache := &authInvalidationCacheStub{ + deleteFn: func(context.Context, string) error { return tc.deleteErr }, + publishFn: func(context.Context, string) error { return tc.publishErr }, + } + worker := NewAuthCacheInvalidationWorker(repo, cache) + worker.processEvent(context.Background(), AuthCacheInvalidationEvent{ID: 9, CacheKey: "hash"}) + require.Equal(t, []int64{9}, repo.retried) + require.Len(t, cache.published, tc.published) + require.NotEmpty(t, repo.retryError) + require.Empty(t, repo.deleted) + require.Equal(t, uint64(1), worker.Health(context.Background()).Failures) + }) + } +} + +func TestAuthCacheInvalidationWorker_RedisSlowIsTimedOut(t *testing.T) { + repo := &authInvalidationRepoStub{} + cache := &authInvalidationCacheStub{deleteFn: func(ctx context.Context, _ string) error { + <-ctx.Done() + return ctx.Err() + }} + worker := NewAuthCacheInvalidationWorker(repo, cache) + started := time.Now() + worker.processEvent(context.Background(), AuthCacheInvalidationEvent{ID: 10, CacheKey: "hash"}) + require.Less(t, time.Since(started), 3*time.Second) + require.Equal(t, []int64{10}, repo.retried) + require.Contains(t, repo.retryError, "deadline") +} + +func TestAuthCacheInvalidationWorker_BoundedBatchAndHealth(t *testing.T) { + oldest := time.Now().Add(-time.Minute) + repo := &authInvalidationRepoStub{stats: AuthCacheInvalidationOutboxStats{ + Pending: 12, OldestCreatedAt: &oldest, MaxAttempts: 4, LastError: "redis down", + }} + worker := NewAuthCacheInvalidationWorker(repo, &authInvalidationCacheStub{}) + require.NoError(t, worker.processBatch(context.Background())) + require.Equal(t, authInvalidationBatchSize, repo.claimLimit) + health := worker.Health(context.Background()) + require.Equal(t, int64(12), health.Pending) + require.Equal(t, 4, health.MaxAttempts) + require.Equal(t, "redis down", health.LastError) + require.GreaterOrEqual(t, health.OldestLag, time.Minute) + require.Equal(t, 35*time.Second, health.HealthySLA) + require.Equal(t, 6*time.Minute, health.RecoverySLA) +} + +func TestAuthCacheInvalidationWorker_ProcessesClaimedBatchConcurrently(t *testing.T) { + events := make([]AuthCacheInvalidationEvent, 32) + for i := range events { + events[i] = AuthCacheInvalidationEvent{ID: int64(i + 1), CacheKey: "hash", Stage: 1} + } + repo := &authInvalidationRepoStub{events: events} + cache := &authInvalidationCacheStub{deleteFn: func(context.Context, string) error { + time.Sleep(100 * time.Millisecond) + return nil + }} + worker := NewAuthCacheInvalidationWorker(repo, cache) + started := time.Now() + require.NoError(t, worker.processBatch(context.Background())) + require.Less(t, time.Since(started), time.Second) + require.Len(t, repo.deleted, 32) +} + +func TestAuthCacheInvalidationWorker_LifecycleIsManagedAndIdempotent(t *testing.T) { + worker := NewAuthCacheInvalidationWorker(&authInvalidationRepoStub{}, &authInvalidationCacheStub{}) + worker.Start() + require.Eventually(t, func() bool { return worker.Health(context.Background()).Running }, time.Second, 10*time.Millisecond) + require.NotPanics(t, func() { worker.Stop(); worker.Stop() }) + require.False(t, worker.Health(context.Background()).Running) +} + +func TestAuthInvalidationRetryDelayIsBoundedAndJittered(t *testing.T) { + for attempt := 1; attempt <= 20; attempt++ { + delay := authInvalidationRetryDelay(attempt) + require.GreaterOrEqual(t, delay, 800*time.Millisecond) + require.LessOrEqual(t, delay, 308*time.Second) + } +} + +func TestAuthCacheInvalidationSubscriber_RetriesInitialFailureAndStops(t *testing.T) { + ready := make(chan struct{}) + var calls int + cache := &authInvalidationCacheStub{subscribeFn: func(ctx context.Context, _ func(string)) error { + calls++ + if calls == 1 { + return errors.New("redis starting") + } + NotifyAuthCacheSubscriptionReady(ctx) + close(ready) + <-ctx.Done() + return ctx.Err() + }} + svc := NewAPIKeyService(nil, nil, nil, nil, nil, cache, nil) + localCache, err := ristretto.NewCache(&ristretto.Config{NumCounters: 10, MaxCost: 1, BufferItems: 64}) + require.NoError(t, err) + defer localCache.Close() + svc.authNegativeCacheL1 = localCache + svc.StartAuthCacheInvalidationSubscriber(context.Background()) + select { + case <-ready: + case <-time.After(2 * time.Second): + t.Fatal("subscriber did not retry") + } + require.Eventually(t, func() bool { return svc.AuthCacheInvalidationSubscriberHealth().Connected }, time.Second, 10*time.Millisecond) + require.Equal(t, uint64(1), svc.AuthCacheInvalidationSubscriberHealth().Failures) + require.NotPanics(t, func() { svc.StopAuthCacheInvalidationSubscriber(); svc.StopAuthCacheInvalidationSubscriber() }) +} + +func TestAuthCacheInvalidationSubscriber_ReconnectsAfterRuntimeDisconnect(t *testing.T) { + ready := make(chan int, 2) + var calls int + cache := &authInvalidationCacheStub{subscribeFn: func(ctx context.Context, _ func(string)) error { + calls++ + NotifyAuthCacheSubscriptionReady(ctx) + ready <- calls + if calls == 1 { + return errors.New("connection dropped") + } + <-ctx.Done() + return ctx.Err() + }} + svc := NewAPIKeyService(nil, nil, nil, nil, nil, cache, nil) + localCache, err := ristretto.NewCache(&ristretto.Config{NumCounters: 10, MaxCost: 1, BufferItems: 64}) + require.NoError(t, err) + defer localCache.Close() + svc.authNegativeCacheL1 = localCache + svc.StartAuthCacheInvalidationSubscriber(context.Background()) + + select { + case call := <-ready: + require.Equal(t, 1, call) + case <-time.After(time.Second): + t.Fatal("initial subscription did not start") + } + select { + case call := <-ready: + require.Equal(t, 2, call) + case <-time.After(2 * time.Second): + t.Fatal("subscriber did not reconnect after runtime disconnect") + } + require.Eventually(t, func() bool { return svc.AuthCacheInvalidationSubscriberHealth().Connected }, time.Second, 10*time.Millisecond) + require.Equal(t, uint64(1), svc.AuthCacheInvalidationSubscriberHealth().Failures) + svc.StopAuthCacheInvalidationSubscriber() +} diff --git a/backend/internal/service/channel_monitor_checker.go b/backend/internal/service/channel_monitor_checker.go index ad4058f9e..910a6e80d 100644 --- a/backend/internal/service/channel_monitor_checker.go +++ b/backend/internal/service/channel_monitor_checker.go @@ -149,11 +149,11 @@ func pingEndpointOrigin(ctx context.Context, endpoint string) *int { return &ms } -// providerAdapter 描述某个 provider 在 challenge 检测中需要的 4 件事: +// providerAdapter 描述某个 provider 在 challenge 检测中需要的几件事: // - 拼出请求路径(含 model 占位) // - 序列化请求体 // - 构造鉴权头 -// - 从响应 JSON 中按 path 提取文本(gjson path) +// - 从响应 JSON 中提取文本(默认按 gjson path;需要时可自定义) // // 加新 provider 只需要在 providerAdapters 里增加一个条目,无需触碰 callProvider / validateProvider。 type providerAdapter struct { @@ -161,6 +161,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* 字符串。 @@ -184,7 +185,7 @@ var providerAdapters = map[string]providerAdapter{ "anthropic-version": monitorAnthropicAPIVersion, } }, - textPath: "content.0.text", + extractText: extractAnthropicMonitorText, }, MonitorProviderGemini: { // Gemini 把 model 名写在 URL path 上:/v1beta/models/{model}:generateContent @@ -293,7 +294,34 @@ func callProvider(ctx context.Context, provider, endpoint, apiKey, model, prompt if provider == MonitorProviderOpenAI && apiMode == MonitorAPIModeResponses { return extractOpenAIResponsesText(respBytes), string(respBytes), status, nil } - 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") } // extractOpenAIResponsesText 聚合 Responses API 的最终 assistant 文本。 diff --git a/backend/internal/service/channel_monitor_checker_body_test.go b/backend/internal/service/channel_monitor_checker_body_test.go index bcf7af0b9..3375ee6ac 100644 --- a/backend/internal/service/channel_monitor_checker_body_test.go +++ b/backend/internal/service/channel_monitor_checker_body_test.go @@ -451,3 +451,50 @@ 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 block after thinking", + body: `{"content":[{"type":"thinking","thinking":""},{"type":"text","text":"2"}]}`, + want: "2", + }, + { + name: "single text block", + body: `{"content":[{"type":"text","text":"2"}]}`, + want: "2", + }, + { + name: "thinking only", + body: `{"content":[{"type":"thinking","thinking":""}]}`, + want: "", + }, + { + name: "multiple text blocks", + body: `{"content":[{"type":"text","text":"answer"},{"type":"tool_use","name":"x"},{"type":"text","text":"2"}]}`, + want: "answer\n2", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := extractAnthropicMonitorText([]byte(tt.body)) + if got != tt.want { + t.Fatalf("extractAnthropicMonitorText() = %q, want %q", got, tt.want) + } + }) + } +} + +func TestValidateChallenge_AnthropicTextAfterThinking(t *testing.T) { + body := []byte(`{"content":[{"type":"thinking","thinking":""},{"type":"text","text":"答案是 2"}]}`) + respText := extractAnthropicMonitorText(body) + + if !validateChallenge(respText, "2") { + t.Fatalf("validateChallenge(%q, %q) = false, want true", respText, "2") + } +} diff --git a/backend/internal/service/error_policy_test.go b/backend/internal/service/error_policy_test.go index 2aa7a421d..147143e8d 100644 --- a/backend/internal/service/error_policy_test.go +++ b/backend/internal/service/error_policy_test.go @@ -209,6 +209,50 @@ func TestCheckErrorPolicy(t *testing.T) { body: []byte(`unauthorized`), expected: ErrorPolicySkipped, }, + { + name: "pool_mode_temp_unschedulable_hit_returns_temp_unscheduled", + account: &Account{ + ID: 9, + Type: AccountTypeAPIKey, + Platform: PlatformOpenAI, + Credentials: map[string]any{ + "pool_mode": true, + "temp_unschedulable_enabled": true, + "temp_unschedulable_rules": []any{ + map[string]any{ + "error_code": float64(http.StatusServiceUnavailable), + "keywords": []any{"unavailable"}, + "duration_minutes": float64(30), + }, + }, + }, + }, + statusCode: http.StatusServiceUnavailable, + body: []byte(`Service temporarily unavailable`), + expected: ErrorPolicyTempUnscheduled, + }, + { + name: "pool_mode_temp_unschedulable_miss_returns_skipped", + account: &Account{ + ID: 10, + Type: AccountTypeAPIKey, + Platform: PlatformOpenAI, + Credentials: map[string]any{ + "pool_mode": true, + "temp_unschedulable_enabled": true, + "temp_unschedulable_rules": []any{ + map[string]any{ + "error_code": float64(http.StatusServiceUnavailable), + "keywords": []any{"maintenance"}, + "duration_minutes": float64(30), + }, + }, + }, + }, + statusCode: http.StatusServiceUnavailable, + body: []byte(`Service temporarily unavailable`), + expected: ErrorPolicySkipped, + }, } for _, tt := range tests { @@ -222,7 +266,7 @@ func TestCheckErrorPolicy(t *testing.T) { } } -func TestHandleUpstreamError_PoolModeCustomErrorCodesOverride(t *testing.T) { +func TestHandleUpstreamError_PoolModePolicies(t *testing.T) { t.Run("pool_mode_without_custom_error_codes_still_skips", func(t *testing.T) { repo := &errorPolicyRepoStub{} svc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil) @@ -262,6 +306,72 @@ func TestHandleUpstreamError_PoolModeCustomErrorCodesOverride(t *testing.T) { require.Equal(t, 1, repo.setErrCalls) require.Equal(t, 0, repo.tempCalls) }) + + t.Run("pool_mode_explicit_temp_rule_stops_scheduling", func(t *testing.T) { + repo := &errorPolicyRepoStub{} + svc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil) + account := &Account{ + ID: 32, + Type: AccountTypeAPIKey, + Platform: PlatformOpenAI, + Credentials: map[string]any{ + "pool_mode": true, + "temp_unschedulable_enabled": true, + "temp_unschedulable_rules": []any{ + map[string]any{ + "error_code": float64(http.StatusServiceUnavailable), + "keywords": []any{"unavailable"}, + "duration_minutes": float64(30), + }, + }, + }, + } + + shouldDisable := svc.HandleUpstreamError( + context.Background(), + account, + http.StatusServiceUnavailable, + http.Header{}, + []byte("Service temporarily unavailable"), + ) + + require.True(t, shouldDisable) + require.Equal(t, 0, repo.setErrCalls) + require.Equal(t, 1, repo.tempCalls) + }) + + t.Run("pool_mode_temp_rule_miss_still_skips", func(t *testing.T) { + repo := &errorPolicyRepoStub{} + svc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil) + account := &Account{ + ID: 33, + Type: AccountTypeAPIKey, + Platform: PlatformOpenAI, + Credentials: map[string]any{ + "pool_mode": true, + "temp_unschedulable_enabled": true, + "temp_unschedulable_rules": []any{ + map[string]any{ + "error_code": float64(http.StatusServiceUnavailable), + "keywords": []any{"maintenance"}, + "duration_minutes": float64(30), + }, + }, + }, + } + + shouldDisable := svc.HandleUpstreamError( + context.Background(), + account, + http.StatusServiceUnavailable, + http.Header{}, + []byte("Service temporarily unavailable"), + ) + + require.False(t, shouldDisable) + require.Equal(t, 0, repo.setErrCalls) + require.Equal(t, 0, repo.tempCalls) + }) } // --------------------------------------------------------------------------- @@ -333,6 +443,9 @@ func TestApplyErrorPolicy(t *testing.T) { Type: AccountTypeOAuth, Platform: PlatformAntigravity, Credentials: map[string]any{ + "model_mapping": map[string]any{ + "claude-sonnet-4-5": "claude-sonnet-4-5", + }, "temp_unschedulable_enabled": true, "temp_unschedulable_rules": []any{ map[string]any{ @@ -362,9 +475,10 @@ func TestApplyErrorPolicy(t *testing.T) { var handleErrorCount int p := antigravityRetryLoopParams{ - ctx: context.Background(), - prefix: "[test]", - account: tt.account, + ctx: context.Background(), + prefix: "[test]", + account: tt.account, + requestedModel: "claude-sonnet-4-5", handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult { handleErrorCount++ return nil @@ -382,6 +496,9 @@ func TestApplyErrorPolicy(t *testing.T) { var switchErr *AntigravityAccountSwitchError require.ErrorAs(t, retErr, &switchErr) require.Equal(t, tt.account.ID, switchErr.OriginalAccountID) + require.Zero(t, repo.tempCalls) + require.Len(t, repo.modelRateLimitCalls, 1) + require.Equal(t, "claude-sonnet-4-5", repo.modelRateLimitCalls[0].scope) } else { require.NoError(t, retErr) } @@ -449,9 +566,10 @@ func TestApplyErrorPolicy_GeminiRateLimitBypassesCustomSkip(t *testing.T) { type errorPolicyRepoStub struct { mockAccountRepoForGemini - tempCalls int - setErrCalls int - lastErrorMsg string + tempCalls int + setErrCalls int + lastErrorMsg string + modelRateLimitCalls []modelNotFoundRateLimitCall } func (r *errorPolicyRepoStub) SetTempUnschedulable(ctx context.Context, id int64, until time.Time, reason string) error { @@ -464,3 +582,12 @@ func (r *errorPolicyRepoStub) SetError(ctx context.Context, id int64, errorMsg s r.lastErrorMsg = errorMsg return nil } + +func (r *errorPolicyRepoStub) SetModelRateLimit(_ context.Context, id int64, scope string, resetAt time.Time, reason ...string) error { + call := modelNotFoundRateLimitCall{accountID: id, scope: scope, resetAt: resetAt} + if len(reason) > 0 { + call.reason = reason[0] + } + r.modelRateLimitCalls = append(r.modelRateLimitCalls, call) + return nil +} diff --git a/backend/internal/service/gateway_model_availability.go b/backend/internal/service/gateway_model_availability.go index f34a59165..bc78a7859 100644 --- a/backend/internal/service/gateway_model_availability.go +++ b/backend/internal/service/gateway_model_availability.go @@ -3,17 +3,20 @@ package service import ( "context" "strings" + + "github.com/Wei-Shaw/sub2api/internal/config" ) // ModelAvailabilityDiagnosis describes whether the requested model can be -// served by any configured account in the group, ignoring transient state -// (rate limits, quota auto-pause, runtime blocks). Handlers use this on the -// "no available accounts" error path to distinguish 404 model_not_found from -// 503 service_unavailable. +// served by any persistently eligible account in the group (active with its +// schedulable setting enabled), ignoring transient state such as rate limits, +// overload, temporary unschedulability, and runtime blocks. Handlers use this +// on the "no available accounts" error path to distinguish 404 +// model_not_found from 503 service_unavailable. type ModelAvailabilityDiagnosis struct { - // HasAccountsInPool is true if the group has at least one schedulable - // account on the queried platform (or, for Anthropic/Gemini, on the - // platform plus mixed-scheduled Antigravity accounts). + // HasAccountsInPool is true if the group has at least one persistently + // eligible account on the queried platform (or, for Anthropic/Gemini, on + // the platform plus mixed-scheduled Antigravity accounts). HasAccountsInPool bool // HasModelSupport is true if at least one account's model mapping admits // the requested model. @@ -33,10 +36,11 @@ type ModelAvailabilityDiagnoser interface { ) ModelAvailabilityDiagnosis } -// DiagnoseModelAvailabilityForPlatform inspects schedulable accounts of the -// given platform and returns whether the requested model is configured to be -// served by any of them. It deliberately ignores schedulability, rate limits, -// quotas, and runtime blocks — those are transient. +// DiagnoseModelAvailabilityForPlatform inspects accounts enabled for scheduling +// by persistent configuration and returns whether the requested model is +// configured to be served by any of them. The dedicated repository query +// bypasses scheduler snapshots and deliberately ignores transient rate-limit, +// overload, temporary-unschedulable, expiry, quota, and runtime-block state. // // Safe to call on the error path: returns {true,true} on any internal failure // or when the inputs preclude meaningful diagnosis (empty model, etc.), so @@ -61,9 +65,30 @@ func (s *GatewayService) DiagnoseModelAvailabilityForPlatform( return ModelAvailabilityDiagnosis{HasAccountsInPool: true, HasModelSupport: true} } - // hasForcePlatform=false so Anthropic/Gemini also surface mixed-scheduled - // Antigravity accounts, matching what selection would consider. - accounts, _, err := s.listSchedulableAccounts(ctx, groupID, platform, false) + if s.accountRepo == nil { + return ModelAvailabilityDiagnosis{HasAccountsInPool: true, HasModelSupport: true} + } + + useMixed := platform == PlatformAnthropic || platform == PlatformGemini + platforms := []string{platform} + if useMixed { + platforms = append(platforms, PlatformAntigravity) + } + + queryGroupID := groupID + includeGrouped := false + if useMixed { + // Preserve the generic scheduler's scope rules: an explicit group wins + // for mixed scheduling, while group-less simple mode scans all accounts. + if groupID == nil && s.cfg != nil && s.cfg.RunMode == config.RunModeSimple { + includeGrouped = true + } + } else if s.cfg != nil && s.cfg.RunMode == config.RunModeSimple { + queryGroupID = nil + includeGrouped = true + } + + accounts, err := s.accountRepo.ListModelAvailabilityCandidates(ctx, queryGroupID, platforms, includeGrouped) if err != nil { // Conservative fallback: pretend everything is fine so the caller // returns 503 (we don't want to flip to 404 just because a lookup @@ -73,6 +98,9 @@ func (s *GatewayService) DiagnoseModelAvailabilityForPlatform( diag := ModelAvailabilityDiagnosis{} for i := range accounts { + if useMixed && accounts[i].Platform == PlatformAntigravity && !accounts[i].IsMixedSchedulingEnabled() { + continue + } diag.HasAccountsInPool = true if s.isModelSupportedByAccountWithContext(ctx, &accounts[i], requestedModel) { diag.HasModelSupport = true diff --git a/backend/internal/service/gateway_model_availability_test.go b/backend/internal/service/gateway_model_availability_test.go index bcca0e5e0..ed32f620e 100644 --- a/backend/internal/service/gateway_model_availability_test.go +++ b/backend/internal/service/gateway_model_availability_test.go @@ -5,6 +5,7 @@ package service import ( "context" "testing" + "time" "github.com/stretchr/testify/require" ) @@ -117,6 +118,7 @@ func TestDiagnoseModelAvailabilityForPlatform_WildcardMappingMatches(t *testing. } func TestDiagnoseModelAvailabilityForPlatform_NoMatchingModel_ReturnsNotFoundSignal(t *testing.T) { + groupID := int64(42) repo := &mockAccountRepoForPlatform{ accounts: []Account{ { @@ -124,6 +126,9 @@ func TestDiagnoseModelAvailabilityForPlatform_NoMatchingModel_ReturnsNotFoundSig Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, + AccountGroups: []AccountGroup{ + {GroupID: groupID}, + }, Credentials: map[string]any{"model_mapping": map[string]any{"gpt-5": "gpt-5"}}, }, { @@ -131,6 +136,9 @@ func TestDiagnoseModelAvailabilityForPlatform_NoMatchingModel_ReturnsNotFoundSig Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, + AccountGroups: []AccountGroup{ + {GroupID: groupID}, + }, Credentials: map[string]any{"model_mapping": map[string]any{"gpt-5-mini": "gpt-5-mini"}}, }, }, @@ -141,12 +149,80 @@ func TestDiagnoseModelAvailabilityForPlatform_NoMatchingModel_ReturnsNotFoundSig } svc := &GatewayService{accountRepo: repo, cfg: testConfig()} - diag := svc.DiagnoseModelAvailabilityForPlatform(context.Background(), nil, "gpt-5.1-codex-mini", PlatformOpenAI) + diag := svc.DiagnoseModelAvailabilityForPlatform(context.Background(), &groupID, "gpt-5.1-codex-mini", PlatformOpenAI) require.True(t, diag.HasAccountsInPool, "group has OpenAI accounts") require.False(t, diag.HasModelSupport, "no account mapping admits the requested model — handler should return 404") } +func TestDiagnoseModelAvailabilityForPlatform_RateLimitedSupportingAccountRemainsConfigured(t *testing.T) { + groupID := int64(42) + cooldownUntil := time.Now().Add(time.Hour) + repo := &mockAccountRepoForPlatform{ + accounts: []Account{ + { + ID: 1, + Platform: PlatformAnthropic, + Status: StatusActive, + Schedulable: true, + RateLimitResetAt: &cooldownUntil, + OverloadUntil: &cooldownUntil, + TempUnschedulableUntil: &cooldownUntil, + AccountGroups: []AccountGroup{{GroupID: groupID}}, + Credentials: map[string]any{ + "model_mapping": map[string]any{"claude-opus-4-8": "claude-opus-4-8"}, + }, + }, + }, + accountsByID: map[int64]*Account{}, + } + require.False(t, repo.accounts[0].IsSchedulable(), "test account must be excluded from normal scheduling while cooling down") + svc := &GatewayService{ + accountRepo: repo, + cfg: testConfig(), + schedulerSnapshot: &SchedulerSnapshotService{}, // diagnosis must bypass the transient-only snapshot + } + + diag := svc.DiagnoseModelAvailabilityForPlatform(context.Background(), &groupID, "claude-opus-4-8", PlatformAnthropic) + + require.True(t, diag.HasAccountsInPool) + require.True(t, diag.HasModelSupport, "a configured model remains supported while every matching account is temporarily cooling down") +} + +func TestOpenAIDiagnoseModelAvailabilityForPlatform_RateLimitedSupportingAccountRemainsConfigured(t *testing.T) { + groupID := int64(43) + cooldownUntil := time.Now().Add(time.Hour) + repo := &mockAccountRepoForPlatform{ + accounts: []Account{ + { + ID: 2, + Platform: PlatformOpenAI, + Status: StatusActive, + Schedulable: true, + RateLimitResetAt: &cooldownUntil, + OverloadUntil: &cooldownUntil, + TempUnschedulableUntil: &cooldownUntil, + AccountGroups: []AccountGroup{{GroupID: groupID}}, + Credentials: map[string]any{ + "model_mapping": map[string]any{"claude-opus-4-8": "claude-opus-4-8"}, + }, + }, + }, + accountsByID: map[int64]*Account{}, + } + require.False(t, repo.accounts[0].IsSchedulable(), "test account must be excluded from normal scheduling while cooling down") + svc := &OpenAIGatewayService{ + accountRepo: repo, + cfg: testConfig(), + schedulerSnapshot: &SchedulerSnapshotService{}, // diagnosis must bypass the transient-only snapshot + } + + diag := svc.DiagnoseModelAvailabilityForPlatform(context.Background(), &groupID, "claude-opus-4-8", PlatformOpenAI) + + require.True(t, diag.HasAccountsInPool) + require.True(t, diag.HasModelSupport, "OpenAI-compatible diagnosis must keep transiently limited supporting accounts in the configured pool") +} + func TestDiagnoseModelAvailabilityForPlatform_WrongPlatformFiltersOut(t *testing.T) { // Group has only Anthropic accounts; user routes to OpenAI gateway. // Diagnosis must NOT see Anthropic accounts (listSchedulableAccounts filters diff --git a/backend/internal/service/gateway_multiplatform_test.go b/backend/internal/service/gateway_multiplatform_test.go index 35a60c812..7b5d967cf 100644 --- a/backend/internal/service/gateway_multiplatform_test.go +++ b/backend/internal/service/gateway_multiplatform_test.go @@ -156,6 +156,34 @@ func (m *mockAccountRepoForPlatform) ListSchedulableUngroupedByPlatform(ctx cont func (m *mockAccountRepoForPlatform) ListSchedulableUngroupedByPlatforms(ctx context.Context, platforms []string) ([]Account, error) { return m.ListSchedulableByPlatforms(ctx, platforms) } +func (m *mockAccountRepoForPlatform) ListModelAvailabilityCandidates(_ context.Context, groupID *int64, platforms []string, includeGrouped bool) ([]Account, error) { + platformSet := make(map[string]struct{}, len(platforms)) + for _, platform := range platforms { + platformSet[platform] = struct{}{} + } + result := make([]Account, 0, len(m.accounts)) + for _, acc := range m.accounts { + if _, ok := platformSet[acc.Platform]; !ok || acc.Status != StatusActive || !acc.Schedulable { + continue + } + if groupID != nil { + inGroup := false + for _, accountGroup := range acc.AccountGroups { + if accountGroup.GroupID == *groupID { + inGroup = true + break + } + } + if !inGroup { + continue + } + } else if !includeGrouped && (len(acc.AccountGroups) > 0 || len(acc.GroupIDs) > 0) { + continue + } + result = append(result, acc) + } + return result, nil +} func (m *mockAccountRepoForPlatform) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error { return nil } diff --git a/backend/internal/service/gateway_non_streaming_response_test.go b/backend/internal/service/gateway_non_streaming_response_test.go index 2416e3e0d..a812a62e4 100644 --- a/backend/internal/service/gateway_non_streaming_response_test.go +++ b/backend/internal/service/gateway_non_streaming_response_test.go @@ -17,8 +17,11 @@ import ( type nonJSONTempUnschedAccountRepo struct { AccountRepository - tempUnschedCalls int - tempReason string + tempUnschedCalls int + tempReason string + modelRateLimitCalls int + modelScope string + modelReason string } func (r *nonJSONTempUnschedAccountRepo) SetTempUnschedulable(_ context.Context, _ int64, _ time.Time, reason string) error { @@ -27,6 +30,15 @@ func (r *nonJSONTempUnschedAccountRepo) SetTempUnschedulable(_ context.Context, return nil } +func (r *nonJSONTempUnschedAccountRepo) SetModelRateLimit(_ context.Context, _ int64, scope string, _ time.Time, reason ...string) error { + r.modelRateLimitCalls++ + r.modelScope = scope + if len(reason) > 0 { + r.modelReason = reason[0] + } + return nil +} + func TestHandleNonStreamingResponse_NonJSON2xxTriggersFailover(t *testing.T) { gin.SetMode(gin.TestMode) rec := httptest.NewRecorder() @@ -131,7 +143,7 @@ func TestHandleNonStreamingResponseAnthropicAPIKeyPassthrough_ValidJSONUnchanged require.JSONEq(t, string(body), rec.Body.String()) } -func TestHandleNonStreamingResponse_NonJSON2xxMatchesTempUnschedulableRule(t *testing.T) { +func TestHandleNonStreamingResponse_NonJSON2xxMatchesModelScopedTempUnschedulableRule(t *testing.T) { gin.SetMode(gin.TestMode) rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) @@ -171,7 +183,9 @@ func TestHandleNonStreamingResponse_NonJSON2xxMatchesTempUnschedulableRule(t *te require.True(t, errors.As(err, &failoverErr)) require.Equal(t, http.StatusBadGateway, failoverErr.StatusCode) require.Equal(t, body, failoverErr.ResponseBody) - require.Equal(t, 1, repo.tempUnschedCalls) - require.Contains(t, repo.tempReason, `"status_code":502`) - require.Contains(t, repo.tempReason, `"matched_keyword":"upstream request failed"`) + require.Zero(t, repo.tempUnschedCalls) + require.Equal(t, 1, repo.modelRateLimitCalls) + require.Equal(t, "claude-sonnet-4-6", repo.modelScope) + require.Contains(t, repo.modelReason, `"status_code":502`) + require.Contains(t, repo.modelReason, `"matched_keyword":"upstream request failed"`) } diff --git a/backend/internal/service/gateway_request.go b/backend/internal/service/gateway_request.go index 1665b1fe4..4ffd5ee8e 100644 --- a/backend/internal/service/gateway_request.go +++ b/backend/internal/service/gateway_request.go @@ -161,6 +161,18 @@ func setGatewayRequestRanges(parsed *ParsedRequest, protocol string, jsonStr str } } +const claudeCodeLongContextModelSuffix = "[1m]" + +// Claude Code treats [1m] as a client-side context selector and normally removes it +// before provider requests. Normalize leaked suffixes, including its duplicated form. +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 +} + // parseGatewayRequestCurrentBody 只做标量和 raw range 轻量解析,不恢复 system/messages 对象图。 func parseGatewayRequestCurrentBody(parsed *ParsedRequest, protocol string) error { if parsed == nil || parsed.Body == nil { @@ -183,6 +195,19 @@ func parseGatewayRequestCurrentBody(parsed *ParsedRequest, protocol string) erro return 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(bodyBytes, "model", normalizedModel) + if err != nil { + return fmt.Errorf("normalize model field: %w", err) + } + parsed.Body.Replace(normalizedBody) + bodyBytes = normalizedBody + jsonStr = *(*string)(unsafe.Pointer(&bodyBytes)) + parsed.Model = normalizedModel + } + } } streamResult := gjson.Get(jsonStr, "stream") diff --git a/backend/internal/service/gateway_request_test.go b/backend/internal/service/gateway_request_test.go index 3172b4d08..d4e010a5e 100644 --- a/backend/internal/service/gateway_request_test.go +++ b/backend/internal/service/gateway_request_test.go @@ -77,6 +77,40 @@ func TestParseGatewayRequest_InvalidStreamType(t *testing.T) { require.Error(t, err) } +func TestParseGatewayRequest_AnthropicNormalizesClaudeCodeLongContextModelSuffix(t *testing.T) { + tests := []struct { + name string + model string + want string + }{ + {name: "lowercase suffix", model: "claude-opus-4-8[1m]", want: "claude-opus-4-8"}, + {name: "uppercase suffix", model: "claude-opus-4-8[1M]", want: "claude-opus-4-8"}, + {name: "duplicated suffix", model: "claude-opus-4-8[1M][1m]", want: "claude-opus-4-8"}, + {name: "suffix in middle", model: "claude-opus-4-8[1m]-preview", want: "claude-opus-4-8[1m]-preview"}, + {name: "suffix only", model: "[1m]", want: "[1m]"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + body := []byte(fmt.Sprintf(`{"model":%q,"system":"test","messages":[{"role":"user","content":"hi"}]}`, tt.model)) + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), domain.PlatformAnthropic) + require.NoError(t, err) + require.Equal(t, tt.want, parsed.Model) + require.Equal(t, tt.want, gjson.GetBytes(parsed.Body.Bytes(), "model").String()) + require.Equal(t, `"test"`, string(parsed.SystemRaw())) + require.NotEmpty(t, parsed.MessagesRaw()) + }) + } +} + +func TestParseGatewayRequest_NonAnthropicPreservesClaudeCodeLongContextModelSuffix(t *testing.T) { + body := []byte(`{"model":"claude-opus-4-8[1m]","input":"hi"}`) + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), "responses") + require.NoError(t, err) + require.Equal(t, "claude-opus-4-8[1m]", parsed.Model) + require.Equal(t, "claude-opus-4-8[1m]", gjson.GetBytes(parsed.Body.Bytes(), "model").String()) +} + func TestParseGatewayRequest_ResponsesInput(t *testing.T) { body := []byte(`{"model":"gpt-5.1","input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"hello"}]}]}`) parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), "responses") diff --git a/backend/internal/service/gemini_chat_completions_compat_service.go b/backend/internal/service/gemini_chat_completions_compat_service.go index 2dc71934c..7db000d26 100644 --- a/backend/internal/service/gemini_chat_completions_compat_service.go +++ b/backend/internal/service/gemini_chat_completions_compat_service.go @@ -143,7 +143,7 @@ func (s *GeminiMessagesCompatService) forwardClaudeBodyAsChatCompletions( return nil, s.writeChatCompletionsError(c, http.StatusBadGateway, "upstream_error", "Upstream request failed after retries: "+safeErr) } - if matched, rebuilt := s.checkErrorPolicyInLoop(ctx, account, resp); matched { + if matched, rebuilt := s.checkErrorPolicyInLoop(ctx, account, resp, mappedModel); matched { resp = rebuilt break } else { @@ -211,7 +211,13 @@ func (s *GeminiMessagesCompatService) forwardClaudeBodyAsChatCompletions( if resp.StatusCode >= 400 { respBody := s.readUpstreamErrorBody(resp) - s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) + policy := ErrorPolicyNone + if s.rateLimitService != nil { + policy = s.rateLimitService.CheckErrorPolicy(ctx, account, resp.StatusCode, respBody, mappedModel) + } + if policy != ErrorPolicyTempUnscheduled { + s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) + } evBody := unwrapIfNeeded(account.Type == AccountTypeOAuth, respBody) if s.shouldFailoverGeminiUpstreamError(resp.StatusCode) { diff --git a/backend/internal/service/gemini_error_policy_test.go b/backend/internal/service/gemini_error_policy_test.go index 84f9a706b..deae912c9 100644 --- a/backend/internal/service/gemini_error_policy_test.go +++ b/backend/internal/service/gemini_error_policy_test.go @@ -199,6 +199,7 @@ func TestGeminiErrorPolicyIntegration(t *testing.T) { expectFailover bool // expect UpstreamFailoverError expectHandleError bool // expect handleGeminiUpstreamError to be called expectShouldFailover bool // for None path, whether shouldFailover triggers + expectModelScope string }{ { name: "custom_codes_matched_429_failover", @@ -252,7 +253,8 @@ func TestGeminiErrorPolicyIntegration(t *testing.T) { statusCode: 503, respBody: []byte(`overloaded`), expectFailover: true, - expectHandleError: true, + expectHandleError: false, + expectModelScope: "gemini-2.5-pro", }, { name: "no_policy_429_failover_via_shouldFailover", @@ -306,17 +308,22 @@ func TestGeminiErrorPolicyIntegration(t *testing.T) { headers := http.Header{} if svc.rateLimitService != nil { - switch svc.rateLimitService.CheckErrorPolicy(ctx, account, statusCode, respBody) { + policy := svc.rateLimitService.CheckErrorPolicy(ctx, account, statusCode, respBody, "gemini-2.5-pro") + switch policy { case ErrorPolicySkipped: // Skipped → return error directly (no handleGeminiUpstreamError, no failover) gotFailover = false handleErrorCalled = false goto verify - case ErrorPolicyMatched, ErrorPolicyTempUnscheduled: + case ErrorPolicyMatched: svc.handleGeminiUpstreamError(ctx, account, statusCode, headers, respBody) handleErrorCalled = true gotFailover = true goto verify + case ErrorPolicyTempUnscheduled: + handleErrorCalled = false + gotFailover = true + goto verify } } @@ -330,6 +337,12 @@ func TestGeminiErrorPolicyIntegration(t *testing.T) { verify: require.Equal(t, tt.expectFailover, gotFailover, "failover mismatch") require.Equal(t, tt.expectHandleError, handleErrorCalled, "handleGeminiUpstreamError call mismatch") + if tt.expectModelScope != "" { + require.Equal(t, 1, repo.setModelRateLimitedCalls) + require.Equal(t, tt.expectModelScope, repo.lastModelScope) + require.Zero(t, repo.setTempCalls) + require.Zero(t, repo.setRateLimitedCalls, "model temp rule must not be widened into an account rate limit") + } if tt.expectShouldFailover { require.True(t, svc.shouldFailoverGeminiUpstreamError(statusCode), @@ -416,9 +429,11 @@ func TestHandleGeminiUpstreamError_GoogleOneCapacityExhaustedUsesTierCooldown(t type geminiErrorPolicyRepo struct { mockAccountRepoForGemini - setErrorCalls int - setRateLimitedCalls int - setTempCalls int + setErrorCalls int + setRateLimitedCalls int + setTempCalls int + setModelRateLimitedCalls int + lastModelScope string } func (r *geminiErrorPolicyRepo) SetError(_ context.Context, _ int64, _ string) error { @@ -435,3 +450,9 @@ func (r *geminiErrorPolicyRepo) SetTempUnschedulable(_ context.Context, _ int64, r.setTempCalls++ return nil } + +func (r *geminiErrorPolicyRepo) SetModelRateLimit(_ context.Context, _ int64, scope string, _ time.Time, _ ...string) error { + r.setModelRateLimitedCalls++ + r.lastModelScope = scope + return nil +} diff --git a/backend/internal/service/gemini_messages_compat_service.go b/backend/internal/service/gemini_messages_compat_service.go index a50a7fc80..2aaba767e 100644 --- a/backend/internal/service/gemini_messages_compat_service.go +++ b/backend/internal/service/gemini_messages_compat_service.go @@ -867,7 +867,7 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex } // 错误策略优先:匹配则跳过重试直接处理。 - if matched, rebuilt := s.checkErrorPolicyInLoop(ctx, account, resp); matched { + if matched, rebuilt := s.checkErrorPolicyInLoop(ctx, account, resp, mappedModel); matched { resp = rebuilt break } else { @@ -937,7 +937,8 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex respBody := s.readUpstreamErrorBody(resp) // 统一错误策略:自定义错误码 + 临时不可调度 if s.rateLimitService != nil { - switch s.rateLimitService.CheckErrorPolicy(ctx, account, resp.StatusCode, respBody) { + policy := s.rateLimitService.CheckErrorPolicy(ctx, account, resp.StatusCode, respBody, mappedModel) + switch policy { case ErrorPolicySkipped: upstreamReqID := resp.Header.Get(requestIDHeader) if upstreamReqID == "" { @@ -945,7 +946,9 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex } return nil, s.writeGeminiMappedError(c, account, http.StatusInternalServerError, upstreamReqID, respBody) case ErrorPolicyMatched, ErrorPolicyTempUnscheduled: - s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) + if policy == ErrorPolicyMatched { + s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) + } upstreamReqID := resp.Header.Get(requestIDHeader) if upstreamReqID == "" { upstreamReqID = resp.Header.Get("x-goog-request-id") @@ -1336,7 +1339,7 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin. } // 错误策略优先:匹配则跳过重试直接处理。 - if matched, rebuilt := s.checkErrorPolicyInLoop(ctx, account, resp); matched { + if matched, rebuilt := s.checkErrorPolicyInLoop(ctx, account, resp, mappedModel); matched { resp = rebuilt break } else { @@ -1445,7 +1448,8 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin. // 统一错误策略:自定义错误码 + 临时不可调度 if s.rateLimitService != nil { - switch s.rateLimitService.CheckErrorPolicy(ctx, account, resp.StatusCode, respBody) { + policy := s.rateLimitService.CheckErrorPolicy(ctx, account, resp.StatusCode, respBody, mappedModel) + switch policy { case ErrorPolicySkipped: respBody = unwrapIfNeeded(isOAuth, respBody) contentType := resp.Header.Get("Content-Type") @@ -1456,7 +1460,9 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin. c.Data(http.StatusInternalServerError, contentType, respBody) return nil, fmt.Errorf("gemini upstream error: %d (skipped by error policy)", resp.StatusCode) case ErrorPolicyMatched, ErrorPolicyTempUnscheduled: - s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) + if policy == ErrorPolicyMatched { + s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) + } evBody := unwrapIfNeeded(isOAuth, respBody) upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(evBody)) upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) @@ -1631,7 +1637,7 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin. // 返回 true 表示策略已匹配(调用者应 break),resp 已重建可直接使用。 // 返回 false 表示 ErrorPolicyNone,resp 已重建,调用者继续走重试逻辑。 func (s *GeminiMessagesCompatService) checkErrorPolicyInLoop( - ctx context.Context, account *Account, resp *http.Response, + ctx context.Context, account *Account, resp *http.Response, mappedModel string, ) (matched bool, rebuilt *http.Response) { if resp.StatusCode < 400 || s.rateLimitService == nil { return false, resp @@ -1643,7 +1649,7 @@ func (s *GeminiMessagesCompatService) checkErrorPolicyInLoop( Header: resp.Header.Clone(), Body: io.NopCloser(bytes.NewReader(body)), } - policy := s.rateLimitService.CheckErrorPolicy(ctx, account, resp.StatusCode, body) + policy := s.rateLimitService.CheckErrorPolicy(ctx, account, resp.StatusCode, body, mappedModel) return policy != ErrorPolicyNone, rebuilt } diff --git a/backend/internal/service/gemini_multiplatform_test.go b/backend/internal/service/gemini_multiplatform_test.go index 7d5ed0ec9..e4c29852e 100644 --- a/backend/internal/service/gemini_multiplatform_test.go +++ b/backend/internal/service/gemini_multiplatform_test.go @@ -147,6 +147,9 @@ func (m *mockAccountRepoForGemini) ListSchedulableUngroupedByPlatform(ctx contex func (m *mockAccountRepoForGemini) ListSchedulableUngroupedByPlatforms(ctx context.Context, platforms []string) ([]Account, error) { return m.ListSchedulableByPlatforms(ctx, platforms) } +func (m *mockAccountRepoForGemini) ListModelAvailabilityCandidates(ctx context.Context, _ *int64, platforms []string, _ bool) ([]Account, error) { + return m.ListSchedulableByPlatforms(ctx, platforms) +} func (m *mockAccountRepoForGemini) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error { return nil } diff --git a/backend/internal/service/grok_media.go b/backend/internal/service/grok_media.go index 5aebc826f..108000470 100644 --- a/backend/internal/service/grok_media.go +++ b/backend/internal/service/grok_media.go @@ -376,6 +376,16 @@ func (s *OpenAIGatewayService) ForwardGrokMedia( if err != nil { return nil, err } + if endpoint == GrokMediaEndpointImagesGenerations || endpoint == GrokMediaEndpointImagesEdits { + if countOpenAIResponseImageOutputsFromJSONBytes(respBody) <= 0 { + setOpsUpstreamError(c, http.StatusBadGateway, "xAI upstream returned no image output", truncateString(string(respBody), 512)) + return nil, &UpstreamFailoverError{ + StatusCode: http.StatusBadGateway, + ResponseBody: respBody, + ResponseHeaders: resp.Header.Clone(), + } + } + } writeGrokMediaResponse(c, resp, respBody, s.responseHeaderFilter) usage := grokMediaUsageFromResponse(endpoint, requestInfo, respBody) return &OpenAIForwardResult{ @@ -598,14 +608,7 @@ func grokMediaUsageFromResponse(endpoint GrokMediaEndpoint, requestInfo GrokMedi meta := grokMediaUsageMetadata{Usage: usage} switch endpoint { case GrokMediaEndpointImagesGenerations, GrokMediaEndpointImagesEdits: - imageCount := countOpenAIResponseImageOutputsFromJSONBytes(responseBody) - if imageCount <= 0 { - imageCount = requestInfo.N - } - if imageCount <= 0 { - imageCount = 1 - } - meta.ImageCount = imageCount + meta.ImageCount = countOpenAIResponseImageOutputsFromJSONBytes(responseBody) meta.ImageSize = requestInfo.SizeTier meta.ImageInputSize = requestInfo.Size meta.ImageOutputSizes = collectOpenAIResponseImageOutputSizesFromJSONBytes(responseBody) diff --git a/backend/internal/service/grok_quota_service.go b/backend/internal/service/grok_quota_service.go index 45e7be83a..db11fe6cc 100644 --- a/backend/internal/service/grok_quota_service.go +++ b/backend/internal/service/grok_quota_service.go @@ -221,6 +221,23 @@ func (s *GrokQuotaService) ProbeBilling(ctx context.Context, accountID int64) (* }) } +// ProbeMediaEligibility refreshes billing state and evaluates the persisted +// account snapshot used by media scheduling. Probe failures remain fail-closed; +// deterministic persisted states such as forbidden or Free are returned as +// normal ineligibility decisions rather than transport errors. +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 { diff --git a/backend/internal/service/grok_quota_service_test.go b/backend/internal/service/grok_quota_service_test.go index 59b19ef27..728277e62 100644 --- a/backend/internal/service/grok_quota_service_test.go +++ b/backend/internal/service/grok_quota_service_test.go @@ -44,6 +44,18 @@ func (r *grokQuotaAccountRepo) UpdateExtra(_ context.Context, id int64, updates r.updates = make(map[int64]map[string]any) } r.updates[id] = updates + if r.mockAccountRepoForPlatform != nil { + account := r.accountsByID[id] + if account == nil { + return nil + } + if account.Extra == nil { + account.Extra = make(map[string]any) + } + for key, value := range updates { + account.Extra[key] = value + } + } return nil } @@ -824,6 +836,53 @@ func TestGrokQuotaServicePartialBilling403PersistsMediaEligibilitySignal(t *test require.Equal(t, "billing_forbidden", reason) } +func TestGrokQuotaServiceProbeMediaEligibility(t *testing.T) { + t.Run("positive paid evidence enables media", func(t *testing.T) { + usagePercent := 10.0 + monthlyLimit := 15_000.0 + account := healthyGrokQuotaOAuthAccount(60) + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + upstream := &grokHybridUpstream{weeklyUsagePercent: &usagePercent, monthlyLimitCents: &monthlyLimit} + svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream, nil) + + eligible, reason, err := svc.ProbeMediaEligibility(context.Background(), account.ID) + + require.NoError(t, err) + require.True(t, eligible) + require.Equal(t, "eligible", reason) + }) + + t.Run("successful empty billing identifies free account", func(t *testing.T) { + account := healthyGrokQuotaOAuthAccount(61) + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), &grokHybridUpstream{}, nil) + + eligible, reason, err := svc.ProbeMediaEligibility(context.Background(), account.ID) + + require.NoError(t, err) + require.False(t, eligible) + require.Equal(t, "billing_free_tier", reason) + }) + + t.Run("forbidden billing is deterministic ineligibility", func(t *testing.T) { + account := healthyGrokQuotaOAuthAccount(62) + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), &grokHybridUpstream{billingStatus: http.StatusForbidden}, nil) + + eligible, reason, err := svc.ProbeMediaEligibility(context.Background(), account.ID) + + require.NoError(t, err) + require.False(t, eligible) + require.Equal(t, "billing_forbidden", reason) + }) +} + func TestPreferBillingObservationStatus(t *testing.T) { t.Parallel() diff --git a/backend/internal/service/invalid_auth_abuse_limiter.go b/backend/internal/service/invalid_auth_abuse_limiter.go new file mode 100644 index 000000000..1ce8ea26a --- /dev/null +++ b/backend/internal/service/invalid_auth_abuse_limiter.go @@ -0,0 +1,276 @@ +package service + +import ( + "sync" + "sync/atomic" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" +) + +const invalidAuthAbuseShardCount = 16 + +type invalidAuthAbuseEntry struct { + failures int + windowStart time.Time + blockedUntil time.Time +} + +type invalidAuthAbuseShard struct { + mu sync.Mutex + entries map[string]*invalidAuthAbuseEntry +} + +type invalidAuthOverflow struct { + mu sync.Mutex + failures int + windowStart time.Time + blockedUntil time.Time +} + +type invalidAuthAbuseLimiter struct { + threshold int + window time.Duration + block time.Duration + capacity int64 + shards [invalidAuthAbuseShardCount]invalidAuthAbuseShard + overflow invalidAuthOverflow + now func() time.Time + + tracked atomic.Int64 + recorded atomic.Uint64 + blocked atomic.Uint64 + rejected atomic.Uint64 + expired atomic.Uint64 + overflowed atomic.Uint64 + globalBlocked atomic.Uint64 + cleanupNext atomic.Int64 + cleanupCursor atomic.Uint32 +} + +type InvalidAuthAbuseHealth struct { + Enabled bool `json:"enabled"` + Tracked int64 `json:"tracked"` + Capacity int64 `json:"capacity"` + Recorded uint64 `json:"recorded"` + Blocks uint64 `json:"blocks"` + Rejected uint64 `json:"rejected"` + Expired uint64 `json:"expired"` + Overflowed uint64 `json:"overflowed"` + GlobalBlocked uint64 `json:"global_blocked"` +} + +func newInvalidAuthAbuseLimiter(cfg *config.Config) *invalidAuthAbuseLimiter { + if cfg == nil || !cfg.APIKeyAuth.InvalidAbuse.Enabled { + return nil + } + c := cfg.APIKeyAuth.InvalidAbuse + if c.Threshold <= 0 || c.WindowSeconds <= 0 || c.BlockSeconds <= 0 || c.Capacity <= 0 { + return nil + } + l := &invalidAuthAbuseLimiter{ + threshold: c.Threshold, + window: time.Duration(c.WindowSeconds) * time.Second, + block: time.Duration(c.BlockSeconds) * time.Second, + capacity: int64(c.Capacity), + now: time.Now, + } + for i := range l.shards { + l.shards[i].entries = make(map[string]*invalidAuthAbuseEntry) + } + return l +} + +func (s *APIKeyService) CheckInvalidAuthAbuse(clientKey string) (time.Duration, bool) { + if s == nil || s.invalidAuthAbuse == nil { + return 0, false + } + return s.invalidAuthAbuse.check(clientKey) +} + +func (s *APIKeyService) RecordInvalidAuthFailure(clientKey string) { + if s == nil || s.invalidAuthAbuse == nil { + return + } + s.invalidAuthAbuse.record(clientKey) +} + +func (s *APIKeyService) InvalidAuthAbuseHealth() InvalidAuthAbuseHealth { + if s == nil || s.invalidAuthAbuse == nil { + return InvalidAuthAbuseHealth{} + } + return s.invalidAuthAbuse.health() +} + +func (l *invalidAuthAbuseLimiter) check(clientKey string) (time.Duration, bool) { + if l == nil || clientKey == "" { + return 0, false + } + now := l.now() + l.maybeCleanupAtCapacity(now) + shard := l.shard(clientKey) + shard.mu.Lock() + entry := shard.entries[clientKey] + if entry != nil && l.entryExpired(entry, now) { + delete(shard.entries, clientKey) + l.tracked.Add(-1) + l.expired.Add(1) + entry = nil + } + if entry != nil && entry.blockedUntil.After(now) { + retry := entry.blockedUntil.Sub(now) + shard.mu.Unlock() + l.rejected.Add(1) + return retry, true + } + shard.mu.Unlock() + + if entry == nil && l.tracked.Load() >= l.capacity { + if retry, blocked := l.checkOverflow(now); blocked { + l.rejected.Add(1) + l.globalBlocked.Add(1) + return retry, true + } + } + return 0, false +} + +func (l *invalidAuthAbuseLimiter) record(clientKey string) { + if l == nil || clientKey == "" { + return + } + l.recorded.Add(1) + now := l.now() + l.maybeCleanupAtCapacity(now) + shard := l.shard(clientKey) + shard.mu.Lock() + entry := shard.entries[clientKey] + if entry != nil && l.entryExpired(entry, now) { + delete(shard.entries, clientKey) + l.tracked.Add(-1) + l.expired.Add(1) + entry = nil + } + if entry == nil { + if !l.reserveEntry() { + shard.mu.Unlock() + l.recordOverflow(now) + return + } + entry = &invalidAuthAbuseEntry{windowStart: now} + shard.entries[clientKey] = entry + } + if entry.blockedUntil.After(now) { + shard.mu.Unlock() + return + } + if entry.windowStart.After(now) || !now.Before(entry.windowStart.Add(l.window)) { + entry.windowStart = now + entry.failures = 0 + } + entry.failures++ + if entry.failures >= l.threshold { + entry.failures = 0 + entry.blockedUntil = now.Add(l.block) + entry.windowStart = entry.blockedUntil + l.blocked.Add(1) + } + shard.mu.Unlock() +} + +func (l *invalidAuthAbuseLimiter) reserveEntry() bool { + for { + current := l.tracked.Load() + if current >= l.capacity { + return false + } + if l.tracked.CompareAndSwap(current, current+1) { + return true + } + } +} + +func (l *invalidAuthAbuseLimiter) maybeCleanupAtCapacity(now time.Time) { + if l.tracked.Load() < l.capacity { + return + } + nowUnixNano := now.UnixNano() + for { + next := l.cleanupNext.Load() + if nowUnixNano < next { + return + } + if l.cleanupNext.CompareAndSwap(next, now.Add(100*time.Millisecond).UnixNano()) { + break + } + } + index := l.cleanupCursor.Add(1) - 1 + shard := &l.shards[index%invalidAuthAbuseShardCount] + shard.mu.Lock() + for key, entry := range shard.entries { + if l.entryExpired(entry, now) { + delete(shard.entries, key) + l.tracked.Add(-1) + l.expired.Add(1) + } + } + shard.mu.Unlock() +} + +func (l *invalidAuthAbuseLimiter) entryExpired(entry *invalidAuthAbuseEntry, now time.Time) bool { + return entry != nil && !entry.blockedUntil.After(now) && !entry.windowStart.After(now) && !now.Before(entry.windowStart.Add(l.window)) +} + +func (l *invalidAuthAbuseLimiter) shard(clientKey string) *invalidAuthAbuseShard { + const fnvOffset32 = uint32(2166136261) + const fnvPrime32 = uint32(16777619) + hash := fnvOffset32 + for i := 0; i < len(clientKey); i++ { + hash ^= uint32(clientKey[i]) + hash *= fnvPrime32 + } + return &l.shards[hash%invalidAuthAbuseShardCount] +} + +func (l *invalidAuthAbuseLimiter) recordOverflow(now time.Time) { + l.overflowed.Add(1) + l.overflow.mu.Lock() + defer l.overflow.mu.Unlock() + if l.overflow.blockedUntil.After(now) { + return + } + if l.overflow.windowStart.IsZero() || !now.Before(l.overflow.windowStart.Add(l.window)) { + l.overflow.windowStart = now + l.overflow.failures = 0 + } + l.overflow.failures++ + if l.overflow.failures >= l.threshold { + l.overflow.failures = 0 + l.overflow.blockedUntil = now.Add(l.block) + l.overflow.windowStart = l.overflow.blockedUntil + l.blocked.Add(1) + } +} + +func (l *invalidAuthAbuseLimiter) checkOverflow(now time.Time) (time.Duration, bool) { + l.overflow.mu.Lock() + defer l.overflow.mu.Unlock() + if l.overflow.blockedUntil.After(now) { + return l.overflow.blockedUntil.Sub(now), true + } + return 0, false +} + +func (l *invalidAuthAbuseLimiter) health() InvalidAuthAbuseHealth { + return InvalidAuthAbuseHealth{ + Enabled: true, + Tracked: l.tracked.Load(), + Capacity: l.capacity, + Recorded: l.recorded.Load(), + Blocks: l.blocked.Load(), + Rejected: l.rejected.Load(), + Expired: l.expired.Load(), + Overflowed: l.overflowed.Load(), + GlobalBlocked: l.globalBlocked.Load(), + } +} diff --git a/backend/internal/service/invalid_auth_abuse_limiter_test.go b/backend/internal/service/invalid_auth_abuse_limiter_test.go new file mode 100644 index 000000000..0d4ac1a45 --- /dev/null +++ b/backend/internal/service/invalid_auth_abuse_limiter_test.go @@ -0,0 +1,100 @@ +//go:build unit + +package service + +import ( + "fmt" + "sync" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/stretchr/testify/require" +) + +func newInvalidAuthLimiterForTest(threshold, capacity int) *invalidAuthAbuseLimiter { + cfg := &config.Config{APIKeyAuth: config.APIKeyAuthCacheConfig{ + InvalidAbuse: config.InvalidAuthAbuseConfig{ + Enabled: true, Threshold: threshold, WindowSeconds: 60, BlockSeconds: 10, Capacity: capacity, + }, + }} + return newInvalidAuthAbuseLimiter(cfg) +} + +func TestInvalidAuthAbuseLimiterBlocksAndExpires(t *testing.T) { + l := newInvalidAuthLimiterForTest(3, 16) + now := time.Date(2026, 7, 17, 0, 0, 0, 0, time.UTC) + l.now = func() time.Time { return now } + + for range 3 { + l.record("203.0.113.1") + } + retry, blocked := l.check("203.0.113.1") + require.True(t, blocked) + require.Equal(t, 10*time.Second, retry) + + now = now.Add(11 * time.Second) + _, blocked = l.check("203.0.113.1") + require.False(t, blocked) + now = now.Add(61 * time.Second) + _, blocked = l.check("203.0.113.1") + require.False(t, blocked) + require.Zero(t, l.health().Tracked) +} + +func TestInvalidAuthAbuseLimiterCapacityUsesBoundedOverflowProtection(t *testing.T) { + l := newInvalidAuthLimiterForTest(2, 2) + now := time.Now() + l.now = func() time.Time { return now } + l.record("198.51.100.1") + l.record("198.51.100.2") + l.record("198.51.100.3") + l.record("198.51.100.4") + + _, blocked := l.check("198.51.100.5") + require.True(t, blocked) + _, trackedBlocked := l.check("198.51.100.1") + require.False(t, trackedBlocked, "global overflow protection should spare existing tracked NATs") + health := l.health() + require.Equal(t, int64(2), health.Tracked) + require.Equal(t, uint64(2), health.Overflowed) + require.Equal(t, uint64(1), health.GlobalBlocked) +} + +func TestInvalidAuthAbuseLimiterConcurrentCapacityIsBounded(t *testing.T) { + const capacity = 64 + l := newInvalidAuthLimiterForTest(1000, capacity) + var wg sync.WaitGroup + for i := range 1000 { + wg.Add(1) + go func(i int) { + defer wg.Done() + l.record(fmt.Sprintf("198.51.100.%d", i)) + }(i) + } + wg.Wait() + health := l.health() + require.LessOrEqual(t, health.Tracked, int64(capacity)) + require.Equal(t, uint64(1000), health.Recorded) + require.Equal(t, uint64(1000-capacity), health.Overflowed) +} + +func TestInvalidAuthAbuseLimiterReclaimsExpiredCapacity(t *testing.T) { + const capacity = 16 + l := newInvalidAuthLimiterForTest(100, capacity) + now := time.Now() + l.now = func() time.Time { return now } + for i := range capacity { + l.record(fmt.Sprintf("source-%d", i)) + } + require.Equal(t, int64(capacity), l.health().Tracked) + + now = now.Add(61 * time.Second) + for i := range invalidAuthAbuseShardCount { + l.check(fmt.Sprintf("new-source-%d", i)) + now = now.Add(101 * time.Millisecond) + } + require.Less(t, l.health().Tracked, int64(capacity)) + l.record("fresh-source") + require.LessOrEqual(t, l.health().Tracked, int64(capacity)) +} diff --git a/backend/internal/service/openai_account_runtime_block_fastpath.go b/backend/internal/service/openai_account_runtime_block_fastpath.go index 7b3df522f..6b68b4e6f 100644 --- a/backend/internal/service/openai_account_runtime_block_fastpath.go +++ b/backend/internal/service/openai_account_runtime_block_fastpath.go @@ -62,17 +62,30 @@ func (s *OpenAIGatewayService) handleOpenAIAccountUpstreamError(ctx context.Cont return false } + if s == nil || account == nil { + return false + } + stateCtx = withTempUnschedulableModel(stateCtx, canonicalModel) + if s.rateLimitService != nil && len(canonicalModel) > 0 && s.rateLimitService.HandleUpstreamModelNotFound(stateCtx, account, canonicalModel[0], statusCode, responseBody) { + return true + } + // Isolate a custom temporary-unschedulable match to the known upstream + // model before entering the generic account error path. This keeps the + // account available to other models and avoids the account runtime blocker. + if s.rateLimitService != nil && statusCode != http.StatusUnauthorized && len(canonicalModel) > 0 && strings.TrimSpace(canonicalModel[0]) != "" && + s.rateLimitService.HandleTempUnschedulable(stateCtx, account, statusCode, responseBody, canonicalModel[0]) { + return true + } if statusCode == http.StatusTooManyRequests { s.markOpenAIOAuth429RateLimited(stateCtx, account, headers, responseBody) } - if s == nil || account == nil || s.rateLimitService == nil { + if s.rateLimitService == nil { return false } - if len(canonicalModel) > 0 && s.rateLimitService.HandleUpstreamModelNotFound(stateCtx, account, canonicalModel[0], statusCode, responseBody) { - return true - } shouldDisable := s.rateLimitService.HandleUpstreamError(stateCtx, account, statusCode, headers, responseBody) - if shouldDisable { + modelTempMatched := statusCode != http.StatusUnauthorized && tempUnschedulableModel(stateCtx, nil) != "" && + len(matchTempUnschedulableRules(account, statusCode, responseBody)) > 0 + if shouldDisable && !modelTempMatched { s.BlockAccountScheduling(account, time.Time{}, "upstream_disable") } if !shouldDisable && account.Platform == PlatformOpenAI && account.Type == AccountTypeAPIKey && shouldCooldownOpenAITransientUpstreamError(statusCode, responseBody) { diff --git a/backend/internal/service/openai_account_runtime_block_fastpath_test.go b/backend/internal/service/openai_account_runtime_block_fastpath_test.go index 7d7cfbda5..ae4bfa91d 100644 --- a/backend/internal/service/openai_account_runtime_block_fastpath_test.go +++ b/backend/internal/service/openai_account_runtime_block_fastpath_test.go @@ -4,6 +4,7 @@ package service import ( "context" + "errors" "net/http" "testing" "time" @@ -83,6 +84,64 @@ func TestOpenAIRuntimeBlocker_IgnoresNonOpenAIFromRateLimitService(t *testing.T) require.False(t, gateway.isOpenAIAccountRuntimeBlocked(account)) } +// 自 #4547(issue 4527 第4点)起,临时不可调度规则命中已知模型时按模型隔离: +// 只封 (账号, 模型) 对,不再账号级一刀切;未知模型仍走账号级兜底 +// (见 TestOpenAITempUnschedulable_UnknownModelKeepsAccountRuntimeBlock)。 +// 池模式规则仍然生效(issue 4470):停止同账号重试并对命中模型设临时封锁。 +func TestOpenAIPoolModeTempRule_StopsSameAccountRetryAndIsolatesBlockToModel(t *testing.T) { + repo := &errorPolicyRepoStub{} + rateLimitService := NewRateLimitService(repo, nil, &config.Config{}, nil, nil) + gateway := &OpenAIGatewayService{ + cfg: &config.Config{}, + rateLimitService: rateLimitService, + } + rateLimitService.SetAccountRuntimeBlocker(gateway) + account := &Account{ + ID: 46, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Credentials: map[string]any{ + "pool_mode": true, + "pool_mode_retry_status_codes": []any{float64(http.StatusServiceUnavailable)}, + "temp_unschedulable_enabled": true, + "temp_unschedulable_rules": []any{ + map[string]any{ + "error_code": float64(http.StatusServiceUnavailable), + "keywords": []any{"unavailable"}, + "duration_minutes": float64(30), + }, + }, + }, + } + body := []byte(`{"error":{"message":"Service temporarily unavailable"}}`) + resp := &http.Response{ + StatusCode: http.StatusServiceUnavailable, + Header: http.Header{}, + } + + failoverErr := gateway.failoverOpenAIUpstreamHTTPError( + context.Background(), + nil, + account, + resp, + body, + "Service temporarily unavailable", + "gpt-5.4", + ) + + require.NotNil(t, failoverErr) + require.False(t, failoverErr.RetryableOnSameAccount) + require.Zero(t, repo.tempCalls) + require.Equal(t, 0, repo.setErrCalls) + require.Equal(t, StatusActive, account.Status) + require.Len(t, repo.modelRateLimitCalls, 1) + require.Equal(t, "gpt-5.4", repo.modelRateLimitCalls[0].scope) + require.False(t, gateway.isOpenAIAccountRuntimeBlocked(account)) + require.False(t, gateway.isOpenAIAccountRequestRuntimeBlocked(account, "gpt-5.5")) +} + func TestOpenAIModelNotFound_DoesNotRuntimeBlockWholeAccount(t *testing.T) { repo := &modelNotFoundAccountRepoStub{} svc := &OpenAIGatewayService{ @@ -105,6 +164,131 @@ func TestOpenAIModelNotFound_DoesNotRuntimeBlockWholeAccount(t *testing.T) { require.Len(t, repo.modelRateLimitCalls, 1) } +func TestOpenAIModelTempUnschedulable_DoesNotRuntimeBlockWholeAccount(t *testing.T) { + repo := &modelNotFoundAccountRepoStub{} + svc := &OpenAIGatewayService{ + rateLimitService: &RateLimitService{accountRepo: repo}, + } + account := openAIModelNotFoundTempAccount() + + shouldDisable := svc.handleOpenAIAccountUpstreamError( + context.Background(), + account, + http.StatusNotFound, + http.Header{}, + []byte(`{"error":{"message":"endpoint not found"}}`), + "gpt-5.4", + ) + + require.True(t, shouldDisable) + require.False(t, svc.isOpenAIAccountRuntimeBlocked(account)) + require.Zero(t, repo.tempCalls) + require.Len(t, repo.modelRateLimitCalls, 1) + require.Equal(t, "gpt-5.4", repo.modelRateLimitCalls[0].scope) +} + +func TestOpenAIModelTempUnschedulable_WriteFailureDoesNotRuntimeBlockWholeAccount(t *testing.T) { + repo := &modelNotFoundAccountRepoStub{modelRateLimitErr: errors.New("write failed")} + svc := &OpenAIGatewayService{ + rateLimitService: &RateLimitService{accountRepo: repo}, + } + account := openAIModelNotFoundTempAccount() + + shouldDisable := svc.handleOpenAIAccountUpstreamError( + context.Background(), + account, + http.StatusNotFound, + http.Header{}, + []byte(`{"error":{"message":"endpoint not found"}}`), + "gpt-5.4", + ) + + require.True(t, shouldDisable) + require.False(t, svc.isOpenAIAccountRuntimeBlocked(account)) + require.Zero(t, repo.tempCalls) + require.Len(t, repo.modelRateLimitCalls, 1) +} + +func TestOpenAIOAuth429_MatchingModelTempRuleAvoidsAccountRuntimeBlock(t *testing.T) { + repo := &modelNotFoundAccountRepoStub{} + svc := &OpenAIGatewayService{ + rateLimitService: &RateLimitService{accountRepo: repo}, + } + account := openAIModelNotFoundTempAccount() + account.Type = AccountTypeOAuth + account.Credentials["temp_unschedulable_rules"] = []any{ + map[string]any{ + "error_code": float64(http.StatusTooManyRequests), + "keywords": []any{"model quota"}, + "duration_minutes": float64(10), + }, + } + + shouldDisable := svc.handleOpenAIAccountUpstreamError( + context.Background(), + account, + http.StatusTooManyRequests, + http.Header{}, + []byte(`{"error":{"message":"model quota exhausted"}}`), + "gpt-5.4", + ) + + require.True(t, shouldDisable) + require.False(t, svc.isOpenAIAccountRuntimeBlocked(account)) + require.Len(t, repo.modelRateLimitCalls, 1) + require.Equal(t, "gpt-5.4", repo.modelRateLimitCalls[0].scope) +} + +func TestOpenAIOAuth429_NonmatchingModelTempRuleKeepsAccountRuntimeBlock(t *testing.T) { + repo := &modelNotFoundAccountRepoStub{} + svc := &OpenAIGatewayService{ + rateLimitService: &RateLimitService{accountRepo: repo}, + } + account := openAIModelNotFoundTempAccount() + account.Type = AccountTypeOAuth + account.Credentials["temp_unschedulable_rules"] = []any{ + map[string]any{ + "error_code": float64(http.StatusTooManyRequests), + "keywords": []any{"different marker"}, + "duration_minutes": float64(10), + }, + } + + shouldDisable := svc.handleOpenAIAccountUpstreamError( + context.Background(), + account, + http.StatusTooManyRequests, + http.Header{}, + []byte(`{"error":{"message":"global rate limit"}}`), + "gpt-5.4", + ) + + require.False(t, shouldDisable) + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) + require.Empty(t, repo.modelRateLimitCalls) +} + +func TestOpenAITempUnschedulable_UnknownModelKeepsAccountRuntimeBlock(t *testing.T) { + repo := &modelNotFoundAccountRepoStub{} + svc := &OpenAIGatewayService{ + rateLimitService: &RateLimitService{accountRepo: repo}, + } + account := openAIModelNotFoundTempAccount() + + shouldDisable := svc.handleOpenAIAccountUpstreamError( + context.Background(), + account, + http.StatusNotFound, + http.Header{}, + []byte(`{"error":{"message":"endpoint not found"}}`), + ) + + require.True(t, shouldDisable) + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) + require.Equal(t, 1, repo.tempCalls) + require.Empty(t, repo.modelRateLimitCalls) +} + func TestOpenAIRuntimeBlock_DoesNotShortenExistingBlock(t *testing.T) { svc := &OpenAIGatewayService{} account := &Account{ID: 46, Platform: PlatformOpenAI, Type: AccountTypeOAuth} diff --git a/backend/internal/service/openai_alpha_search.go b/backend/internal/service/openai_alpha_search.go index f11a21958..488db038e 100644 --- a/backend/internal/service/openai_alpha_search.go +++ b/backend/internal/service/openai_alpha_search.go @@ -95,13 +95,14 @@ func (s *OpenAIGatewayService) ForwardAlphaSearch(ctx context.Context, c *gin.Co // 被永久标记为 error;历史导入且缺少 auth_mode 标记的 at- token 也会 // 漏过 PAT 类型判断。这里仍允许本次请求换号,但不修改任何账号状态; // 真正的凭据失效由普通 Responses 请求或 whoami 校验判定。 + shouldDisable := false if shouldApplyOpenAIAlphaSearchAccountErrorSideEffects(resp.StatusCode) { - s.handleFailoverSideEffects(ctx, resp, account, respBody, upstreamModel) + shouldDisable = s.handleFailoverSideEffects(ctx, resp, account, respBody, upstreamModel) } return nil, &UpstreamFailoverError{ StatusCode: resp.StatusCode, ResponseBody: respBody, - RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), + RetryableOnSameAccount: !shouldDisable && account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), } } } @@ -169,13 +170,14 @@ func (s *OpenAIGatewayService) forwardAlphaSearchViaResponsesWebSearch( if s.shouldFailoverOpenAIUpstreamResponse(resp.StatusCode, upstreamMessage, respBody) { resp.Body = io.NopCloser(bytes.NewReader(respBody)) // 仍按 alpha/search 工具请求处理:PAT 的工具链路失败不能直接永久置错。 + shouldDisable := false if shouldApplyOpenAIAlphaSearchAccountErrorSideEffects(resp.StatusCode) { - s.handleFailoverSideEffects(ctx, resp, account, respBody, upstreamModel) + shouldDisable = s.handleFailoverSideEffects(ctx, resp, account, respBody, upstreamModel) } return nil, &UpstreamFailoverError{ StatusCode: resp.StatusCode, ResponseBody: respBody, - RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), + RetryableOnSameAccount: !shouldDisable && account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), } } } diff --git a/backend/internal/service/openai_embeddings.go b/backend/internal/service/openai_embeddings.go index fa821d8f3..c0c590029 100644 --- a/backend/internal/service/openai_embeddings.go +++ b/backend/internal/service/openai_embeddings.go @@ -132,11 +132,11 @@ func (s *OpenAIGatewayService) ForwardEmbeddings( Message: upstreamMsg, Detail: upstreamDetail, }) - s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody, upstreamModel) + shouldDisable := s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody, upstreamModel) return nil, &UpstreamFailoverError{ StatusCode: resp.StatusCode, ResponseBody: respBody, - RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), + RetryableOnSameAccount: !shouldDisable && account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), } } writeOpenAIEmbeddingsUpstreamResponse(c, resp, respBody, s.responseHeaderFilter) diff --git a/backend/internal/service/openai_gateway_cc_pipeline.go b/backend/internal/service/openai_gateway_cc_pipeline.go index 52658b272..533f2ad38 100644 --- a/backend/internal/service/openai_gateway_cc_pipeline.go +++ b/backend/internal/service/openai_gateway_cc_pipeline.go @@ -112,15 +112,16 @@ func (s *OpenAIGatewayService) failoverOpenAIUpstreamHTTPError( Message: upstreamMsg, Detail: upstreamDetail, }) + shouldDisable := false if account.Platform != PlatformGrok { - s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody, upstreamModel) + shouldDisable = s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody, upstreamModel) } return newOpenAIUpstreamFailoverError( resp.StatusCode, resp.Header, respBody, upstreamMsg, - account.IsPoolMode() && (account.IsPoolModeRetryableStatus(resp.StatusCode) || isOpenAITransientProcessingError(resp.StatusCode, upstreamMsg, respBody)), + !shouldDisable && account.IsPoolMode() && (account.IsPoolModeRetryableStatus(resp.StatusCode) || isOpenAITransientProcessingError(resp.StatusCode, upstreamMsg, respBody)), ) } diff --git a/backend/internal/service/openai_gateway_chat_completions.go b/backend/internal/service/openai_gateway_chat_completions.go index fe24b3db3..cccec4605 100644 --- a/backend/internal/service/openai_gateway_chat_completions.go +++ b/backend/internal/service/openai_gateway_chat_completions.go @@ -779,9 +779,9 @@ func (s *OpenAIGatewayService) handleChatStreamingResponse( handleScanErr := func(err error) { if err != nil && !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded) { - logger.L().Warn("openai chat_completions stream: read error", + logger.FromContext(c.Request.Context()).Warn("openai chat_completions stream: read error", zap.Error(err), - zap.String("request_id", requestID), + zap.String("upstream_request_id", requestID), ) } } @@ -820,7 +820,10 @@ func (s *OpenAIGatewayService) handleChatStreamingResponse( } if err := scanner.Err(); err != nil { handleScanErr(err) - return resultWithUsage(), fmt.Errorf("stream usage incomplete: %w", err) + if clientDisconnected || errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return resultWithUsage(), fmt.Errorf("stream usage incomplete: %w", err) + } + return resultWithUsage(), newOpenAIUpstreamStreamReadError(err) } if frame, ok := parser.Finish(); ok { if strings.TrimSpace(frame.Data) == "[DONE]" { @@ -892,7 +895,10 @@ func (s *OpenAIGatewayService) handleChatStreamingResponse( } if ev.err != nil { handleScanErr(ev.err) - return resultWithUsage(), fmt.Errorf("stream usage incomplete: %w", ev.err) + if clientDisconnected || errors.Is(ev.err, context.Canceled) || errors.Is(ev.err, context.DeadlineExceeded) { + return resultWithUsage(), fmt.Errorf("stream usage incomplete: %w", ev.err) + } + return resultWithUsage(), newOpenAIUpstreamStreamReadError(ev.err) } lastDataAt = time.Now() line := ev.line diff --git a/backend/internal/service/openai_gateway_chat_completions_test.go b/backend/internal/service/openai_gateway_chat_completions_test.go index 5186598a7..b307aa4ed 100644 --- a/backend/internal/service/openai_gateway_chat_completions_test.go +++ b/backend/internal/service/openai_gateway_chat_completions_test.go @@ -32,6 +32,63 @@ func (w *openAIChatFailingWriter) Write(p []byte) (int, error) { return w.ResponseWriter.Write(p) } +type openAIChatStreamReadErrorCloser struct { + payload []byte + err error + sent bool +} + +func (r *openAIChatStreamReadErrorCloser) Read(p []byte) (int, error) { + if !r.sent { + r.sent = true + return copy(p, r.payload), nil + } + return 0, r.err +} + +func (r *openAIChatStreamReadErrorCloser) Close() error { return nil } + +func TestHandleChatStreamingResponse_ClassifiesHTTP2ReadError(t *testing.T) { + gin.SetMode(gin.TestMode) + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + resp := &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"text/event-stream"}, + "x-request-id": []string{"upstream-rid"}, + }, + Body: &openAIChatStreamReadErrorCloser{ + payload: []byte("data: {\"type\":\"response.output_text.delta\",\"delta\":\"partial\"}\n\n"), + err: errors.New("stream error: stream ID 5; INTERNAL_ERROR; received from peer"), + }, + } + svc := &OpenAIGatewayService{cfg: &config.Config{}} + + result, err := svc.handleChatStreamingResponse( + resp, + c, + &Account{ID: 1, Name: "openai-oauth", Platform: PlatformOpenAI}, + "gpt-5.6-sol", + "gpt-5.6-sol", + "gpt-5.6-sol", + time.Now(), + 0, + ) + + require.Error(t, err) + require.NotNil(t, result) + require.True(t, c.Writer.Written(), "partial output must make replay unsafe") + code, message, ok := OpenAIUpstreamStreamReadErrorDetails(err) + require.True(t, ok) + require.Equal(t, OpenAIUpstreamHTTP2StreamErrorCode, code) + require.Equal(t, "Upstream HTTP/2 stream failed", message) + require.NotContains(t, message, "stream ID") + require.NotContains(t, message, "INTERNAL_ERROR") +} + func TestNormalizeResponsesRequestServiceTier(t *testing.T) { t.Parallel() diff --git a/backend/internal/service/openai_gateway_forward.go b/backend/internal/service/openai_gateway_forward.go index 16c034f19..93697be77 100644 --- a/backend/internal/service/openai_gateway_forward.go +++ b/backend/internal/service/openai_gateway_forward.go @@ -877,13 +877,13 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco Detail: upstreamDetail, }) - s.handleFailoverSideEffects(ctx, resp, account, respBody, upstreamModel) + shouldDisable := s.handleFailoverSideEffects(ctx, resp, account, respBody, upstreamModel) return nil, newOpenAIUpstreamFailoverError( resp.StatusCode, resp.Header, respBody, upstreamMsg, - account.IsPoolMode() && (account.IsPoolModeRetryableStatus(resp.StatusCode) || isOpenAITransientProcessingError(resp.StatusCode, upstreamMsg, respBody)), + !shouldDisable && account.IsPoolMode() && (account.IsPoolModeRetryableStatus(resp.StatusCode) || isOpenAITransientProcessingError(resp.StatusCode, upstreamMsg, respBody)), ) } return s.handleErrorResponse(ctx, resp, c, account, body, billingModel) diff --git a/backend/internal/service/openai_gateway_grok_cache.go b/backend/internal/service/openai_gateway_grok_cache.go index 9b6436e43..d9e34cbb5 100644 --- a/backend/internal/service/openai_gateway_grok_cache.go +++ b/backend/internal/service/openai_gateway_grok_cache.go @@ -296,6 +296,13 @@ func appendMissingGrokFreeCacheNativeTools(body []byte) ([]byte, error) { if !hasFunction { return body, nil } + // Only complement missing native search tools when the request already contains + // at least one search tool (native or function-form). Pure client function tools + // (e.g. view_image) must not trigger injection to avoid biasing model tool + // selection (#4486). + if !present["web_search"] && !present["x_search"] { + return body, nil + } for _, toolType := range []string{"web_search", "x_search"} { if present[toolType] { continue diff --git a/backend/internal/service/openai_gateway_grok_cache_test.go b/backend/internal/service/openai_gateway_grok_cache_test.go index c31ee7bc3..73b1ea5ec 100644 --- a/backend/internal/service/openai_gateway_grok_cache_test.go +++ b/backend/internal/service/openai_gateway_grok_cache_test.go @@ -211,6 +211,7 @@ func TestApplyGrokCacheIdentityAppendsNativeToolsToResponseFunctions(t *testing. for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { + // Pure client function tools without search → no native injection (#4486). intentBody := []byte(`{"model":"grok","tools":[{"type":"function","name":"lookup","description":"look up a value","parameters":{"type":"object"}},{"type":"function","name":"save","parameters":{"type":"object"}}]` + tt.toolChoiceJSON + `}`) body, err := applyGrokResponsesCacheIdentity(intentBody, intentBody, "isolated-id", true) require.NoError(t, err) @@ -219,28 +220,35 @@ func TestApplyGrokCacheIdentityAppendsNativeToolsToResponseFunctions(t *testing. require.NoError(t, err) require.Equal(t, "isolated-id", gjson.GetBytes(body, "prompt_cache_key").String()) tools := gjson.GetBytes(body, "tools").Array() - require.Len(t, tools, 4) + require.Len(t, tools, 2, "pure client functions should not get native search injected") require.Equal(t, "function", tools[0].Get("type").String()) require.Equal(t, "lookup", tools[0].Get("name").String()) require.Equal(t, "function", tools[1].Get("type").String()) require.Equal(t, "save", tools[1].Get("name").String()) - require.Equal(t, "web_search", tools[2].Get("type").String()) - require.Equal(t, "x_search", tools[3].Get("type").String()) require.Equal(t, tt.wantChoice, gjson.GetBytes(body, "tool_choice").Exists()) - if tt.wantChoice { - require.Equal(t, "auto", gjson.GetBytes(body, "tool_choice").String()) - } - - second, err := applyGrokResponsesCacheIdentity(body, intentBody, "isolated-id", true) - require.NoError(t, err) - second, err = applyGrokFreeMessagesFunctionToolCacheRoute(second, intentBody, account, "isolated-id") - require.NoError(t, err) - require.JSONEq(t, string(body), string(second), "native tools must not be duplicated") - require.Len(t, gjson.GetBytes(second, "tools").Array(), 4) }) } } +func TestApplyGrokCacheIdentityAppendsNativeToolsWhenSearchPresent(t *testing.T) { + account := healthyGrokOAuthGatewayTestAccount(901, "access-token") + account.Credentials["subscription_tier"] = " FREE " + + // Function tools INCLUDING web_search → convert + complement with x_search. + intentBody := []byte(`{"model":"grok","tools":[{"type":"function","name":"lookup","description":"look up a value","parameters":{"type":"object"}},{"type":"function","name":"web_search","description":"search","parameters":{"type":"object"}}]}`) + body, err := applyGrokResponsesCacheIdentity(intentBody, intentBody, "isolated-id", true) + require.NoError(t, err) + body, err = applyGrokFreeMessagesFunctionToolCacheRoute(body, intentBody, account, "isolated-id") + require.NoError(t, err) + + tools := gjson.GetBytes(body, "tools").Array() + require.Len(t, tools, 3, "lookup(function) + web_search(native) + x_search(native)") + require.Equal(t, "function", tools[0].Get("type").String()) + require.Equal(t, "lookup", tools[0].Get("name").String()) + require.Equal(t, "web_search", tools[1].Get("type").String()) + require.Equal(t, "x_search", tools[2].Get("type").String()) +} + func TestApplyGrokCacheIdentityRequiresPatchedFunctionTools(t *testing.T) { account := healthyGrokOAuthGatewayTestAccount(902, "access-token") account.Credentials["subscription_tier"] = "free" @@ -272,7 +280,9 @@ func TestApplyGrokCacheIdentityRequiresPatchedFunctionTools(t *testing.T) { } func TestGrokFreeMessagesFunctionToolCacheRouteRequiresKnownFreeTier(t *testing.T) { - intentBody := []byte(`{"model":"grok","tools":[{"type":"function","name":"lookup"}],"tool_choice":"auto"}`) + // Include web_search as function to trigger native tool injection (pure client + // functions no longer trigger injection after #4486 fix). + intentBody := []byte(`{"model":"grok","tools":[{"type":"function","name":"lookup"},{"type":"function","name":"web_search"}],"tool_choice":"auto"}`) tests := []struct { name string account *Account @@ -382,7 +392,7 @@ func TestGrokFreeMessagesFunctionToolCacheRouteRequiresKnownFreeTier(t *testing. require.Equal(t, "x_search", tools[2].Get("type").String()) return } - require.Len(t, tools, 1) + require.Len(t, tools, 2, "non-free accounts should not get native search injected") }) } } diff --git a/backend/internal/service/openai_gateway_grok_cache_tool_test.go b/backend/internal/service/openai_gateway_grok_cache_tool_test.go new file mode 100644 index 000000000..77dba26ce --- /dev/null +++ b/backend/internal/service/openai_gateway_grok_cache_tool_test.go @@ -0,0 +1,87 @@ +package service + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +func TestAppendMissingGrokFreeCacheNativeTools_PureClientFunctionNoInject(t *testing.T) { + body := []byte(`{ + "model": "grok-4.5", + "tools": [ + {"type":"function","name":"view_image","description":"View image","parameters":{"type":"object","properties":{"path":{"type":"string"}},"required":["path"]}} + ], + "tool_choice": "auto" + }`) + + result, err := appendMissingGrokFreeCacheNativeTools(body) + require.NoError(t, err) + + tools := gjson.GetBytes(result, "tools").Array() + for _, tool := range tools { + toolType := tool.Get("type").String() + assert.NotEqual(t, "web_search", toolType, "should not inject web_search for pure client functions") + assert.NotEqual(t, "x_search", toolType, "should not inject x_search for pure client functions") + } +} + +func TestAppendMissingGrokFreeCacheNativeTools_FunctionPlusWebSearchInjects(t *testing.T) { + body := []byte(`{ + "model": "grok-4.5", + "tools": [ + {"type":"function","name":"view_image","description":"View","parameters":{"type":"object","properties":{"path":{"type":"string"}},"required":["path"]}}, + {"type":"function","name":"web_search","description":"Search","parameters":{"type":"object","properties":{"query":{"type":"string"}},"required":["query"]}} + ] + }`) + + result, err := appendMissingGrokFreeCacheNativeTools(body) + require.NoError(t, err) + + tools := gjson.GetBytes(result, "tools").Array() + types := make(map[string]bool) + for _, tool := range tools { + types[tool.Get("type").String()] = true + } + assert.True(t, types["web_search"], "web_search should be present (converted from function)") + assert.True(t, types["x_search"], "x_search should be injected when web_search is present alongside client functions") +} + +func TestAppendMissingGrokFreeCacheNativeTools_NativeSearchAlreadyPresent(t *testing.T) { + body := []byte(`{ + "model": "grok-4.5", + "tools": [ + {"type":"function","name":"view_image","description":"View","parameters":{"type":"object","properties":{"path":{"type":"string"}},"required":["path"]}}, + {"type":"web_search"} + ] + }`) + + result, err := appendMissingGrokFreeCacheNativeTools(body) + require.NoError(t, err) + + tools := gjson.GetBytes(result, "tools").Array() + types := make(map[string]bool) + for _, tool := range tools { + types[tool.Get("type").String()] = true + } + assert.True(t, types["web_search"]) + assert.True(t, types["x_search"], "x_search should be injected when web_search is already present") +} + +func TestAppendMissingGrokFreeCacheNativeTools_MultipleFunctionsNoSearch(t *testing.T) { + body := []byte(`{ + "model": "grok-4.5", + "tools": [ + {"type":"function","name":"view_image","description":"View","parameters":{"type":"object","properties":{"path":{"type":"string"}},"required":["path"]}}, + {"type":"function","name":"read_file","description":"Read","parameters":{"type":"object","properties":{"path":{"type":"string"}},"required":["path"]}} + ] + }`) + + result, err := appendMissingGrokFreeCacheNativeTools(body) + require.NoError(t, err) + + tools := gjson.GetBytes(result, "tools").Array() + require.Len(t, tools, 2, "no tools should be injected for pure client functions") +} diff --git a/backend/internal/service/openai_gateway_grok_test.go b/backend/internal/service/openai_gateway_grok_test.go index 9e91eb82d..163ab4e5a 100644 --- a/backend/internal/service/openai_gateway_grok_test.go +++ b/backend/internal/service/openai_gateway_grok_test.go @@ -551,7 +551,7 @@ func TestForwardGrokMediaImagesGenerationNormalizesImagineAlias(t *testing.T) { "Content-Type": []string{"application/json"}, "Xai-Request-Id": []string{"xai-image-req"}, }, - Body: io.NopCloser(strings.NewReader(`{"data":[]}`)), + Body: io.NopCloser(strings.NewReader(`{"data":[{"url":"https://images.test/cat.png"}]}`)), }} svc := &OpenAIGatewayService{httpUpstream: upstream} @@ -565,7 +565,7 @@ func TestForwardGrokMediaImagesGenerationNormalizesImagineAlias(t *testing.T) { require.NotEqual(t, grokUpstreamUserAgent, upstream.lastReq.Header.Get("User-Agent")) require.JSONEq(t, `{"model":"grok-imagine-image-quality","prompt":"draw a cat"}`, string(upstream.lastBody)) require.Equal(t, http.StatusOK, recorder.Code) - require.JSONEq(t, `{"data":[]}`, recorder.Body.String()) + require.JSONEq(t, `{"data":[{"url":"https://images.test/cat.png"}]}`, recorder.Body.String()) require.Equal(t, "xai-image-req", result.RequestID) require.Equal(t, "grok-imagine-image-quality", result.Model) require.Equal(t, "grok-imagine-image-quality", result.BillingModel) @@ -597,7 +597,7 @@ func TestForwardGrokMediaAppliesAccountModelMappingAfterEndpointNormalization(t wantRequestModel: "grok-imagine-image-quality", wantUpstream: "vendor-image-model", wantBody: `{"model":"vendor-image-model","prompt":"draw a cat"}`, - responseBody: `{"data":[]}`, + responseBody: `{"data":[{"url":"https://images.test/mapped.png"}]}`, }, { name: "video generation maps text-only fallback model", @@ -630,7 +630,7 @@ func TestForwardGrokMediaAppliesAccountModelMappingAfterEndpointNormalization(t wantRequestModel: "grok-imagine-image-quality", wantUpstream: "vendor-image-model", wantBody: `{"model":"vendor-image-model","prompt":"draw"}`, - responseBody: `{"data":[]}`, + responseBody: `{"data":[{"url":"https://images.test/mapped.png"}]}`, }, { name: "whitespace mapping target safely preserves normalized model", @@ -641,7 +641,7 @@ func TestForwardGrokMediaAppliesAccountModelMappingAfterEndpointNormalization(t wantRequestModel: "grok-imagine-image-quality", wantUpstream: "grok-imagine-image-quality", wantBody: `{"model":"grok-imagine-image-quality","prompt":"draw"}`, - responseBody: `{"data":[]}`, + responseBody: `{"data":[{"url":"https://images.test/mapped.png"}]}`, }, } @@ -682,6 +682,43 @@ func TestForwardGrokMediaAppliesAccountModelMappingAfterEndpointNormalization(t } } +func TestForwardGrokMediaImagesGenerationRejectsEmptySuccessfulResponse(t *testing.T) { + t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true") + gin.SetMode(gin.TestMode) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + body := []byte(`{"model":"grok-imagine-image","prompt":"draw a cat"}`) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + + account := &Account{ + ID: 66, + Name: "grok", + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "api-key", + "base_url": "https://xai.test/v1", + }, + } + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"data":[]}`)), + }} + svc := &OpenAIGatewayService{httpUpstream: upstream} + + result, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointImagesGenerations, "", body, "application/json") + require.Nil(t, result) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, http.StatusBadGateway, failoverErr.StatusCode) + require.JSONEq(t, `{"data":[]}`, string(failoverErr.ResponseBody)) + require.Empty(t, recorder.Body.String()) +} + func TestForwardGrokMediaImagesGenerationStripsUnsupportedSize(t *testing.T) { t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true") gin.SetMode(gin.TestMode) @@ -709,7 +746,7 @@ func TestForwardGrokMediaImagesGenerationStripsUnsupportedSize(t *testing.T) { Header: http.Header{ "Content-Type": []string{"application/json"}, }, - Body: io.NopCloser(strings.NewReader(`{"data":[]}`)), + Body: io.NopCloser(strings.NewReader(`{"data":[{"url":"https://images.test/cat.png"}]}`)), }} svc := &OpenAIGatewayService{httpUpstream: upstream} @@ -759,7 +796,7 @@ func TestForwardGrokMediaImagesEditMultipartConvertsToJSON(t *testing.T) { Header: http.Header{ "Content-Type": []string{"application/json"}, }, - Body: io.NopCloser(strings.NewReader(`{"data":[]}`)), + Body: io.NopCloser(strings.NewReader(`{"data":[{"url":"https://images.test/edited.png"}]}`)), }} svc := &OpenAIGatewayService{httpUpstream: upstream} @@ -1723,7 +1760,7 @@ func TestForwardAsAnthropicForGrokFunctionToolUsesCacheCapableMixedRoute(t *test body := []byte(`{ "model":"grok","max_tokens":32,"stream":false, "messages":[{"role":"user","content":"look up alpha"}], - "tools":[{"name":"lookup","description":"look up a key","input_schema":{"type":"object","properties":{"key":{"type":"string"}},"required":["key"]}}], + "tools":[{"name":"lookup","description":"look up a key","input_schema":{"type":"object","properties":{"key":{"type":"string"}},"required":["key"]}},{"name":"web_search","description":"search the web","input_schema":{"type":"object","properties":{"query":{"type":"string"}},"required":["query"]}}], "tool_choice":{"type":"auto"} }`) c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body)) diff --git a/backend/internal/service/openai_gateway_model_availability.go b/backend/internal/service/openai_gateway_model_availability.go index 1bbc3ee7c..a665052e6 100644 --- a/backend/internal/service/openai_gateway_model_availability.go +++ b/backend/internal/service/openai_gateway_model_availability.go @@ -3,13 +3,16 @@ package service import ( "context" "strings" + + "github.com/Wei-Shaw/sub2api/internal/config" ) // DiagnoseModelAvailabilityForPlatform reports whether the requested model -// is configured to be served by any OpenAI-compatible account in the group -// for the given platform (e.g. PlatformOpenAI, PlatformGrok). The platform -// scopes the candidate pool so distinct OpenAI-compatible platforms do not -// cross-contaminate diagnosis results. +// is configured to be served by any persistently eligible OpenAI-compatible +// account in the group for the given platform (e.g. PlatformOpenAI, +// PlatformGrok). The platform scopes the candidate pool so distinct +// OpenAI-compatible platforms do not cross-contaminate diagnosis results. +// The query bypasses scheduler snapshots and ignores transient runtime state. // // Safe to call on the error path: returns {true,true} on any internal // failure or when the inputs preclude meaningful diagnosis (empty model, @@ -27,8 +30,23 @@ func (s *OpenAIGatewayService) DiagnoseModelAvailabilityForPlatform( if requestedModel == "" { return ModelAvailabilityDiagnosis{HasAccountsInPool: true, HasModelSupport: true} } + if s.accountRepo == nil { + return ModelAvailabilityDiagnosis{HasAccountsInPool: true, HasModelSupport: true} + } - accounts, err := s.listSchedulableAccounts(ctx, groupID, platform) + platform = normalizeOpenAICompatiblePlatform(platform) + queryGroupID := groupID + includeGrouped := false + if s.cfg != nil && s.cfg.RunMode == config.RunModeSimple { + queryGroupID = nil + includeGrouped = true + } + accounts, err := s.accountRepo.ListModelAvailabilityCandidates( + ctx, + queryGroupID, + []string{platform}, + includeGrouped, + ) if err != nil { // Conservative fallback so the caller keeps returning 503; we do not // want a transient lookup failure to flip into 404 model_not_found. diff --git a/backend/internal/service/openai_gateway_passthrough.go b/backend/internal/service/openai_gateway_passthrough.go index a17d050af..1fb281b4a 100644 --- a/backend/internal/service/openai_gateway_passthrough.go +++ b/backend/internal/service/openai_gateway_passthrough.go @@ -580,7 +580,7 @@ func (s *OpenAIGatewayService) handleFailoverErrorResponsePassthrough( logOpenAIInstructionsRequiredDebug(ctx, c, account, resp.StatusCode, upstreamMsg, requestBody, body) reqModel, _, _ := extractOpenAIRequestMetaFromBody(requestBody) canonicalModel := canonicalOpenAIAccountSchedulingModel(account, reqModel) - _ = s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, body, canonicalModel) + shouldDisable := s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, body, canonicalModel) appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ Platform: account.Platform, AccountID: account.ID, @@ -598,7 +598,7 @@ func (s *OpenAIGatewayService) handleFailoverErrorResponsePassthrough( resp.Header, body, upstreamMsg, - account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), + !shouldDisable && account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), ) } diff --git a/backend/internal/service/openai_gateway_upstream_errors.go b/backend/internal/service/openai_gateway_upstream_errors.go index a21775308..ff97e43fc 100644 --- a/backend/internal/service/openai_gateway_upstream_errors.go +++ b/backend/internal/service/openai_gateway_upstream_errors.go @@ -305,12 +305,11 @@ func (s *OpenAIGatewayService) readUpstreamErrorBody(resp *http.Response) []byte return body } -func (s *OpenAIGatewayService) handleFailoverSideEffects(ctx context.Context, resp *http.Response, account *Account, responseBody []byte, canonicalModel ...string) { +func (s *OpenAIGatewayService) handleFailoverSideEffects(ctx context.Context, resp *http.Response, account *Account, responseBody []byte, canonicalModel ...string) bool { if len(canonicalModel) > 0 { - s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, responseBody, canonicalModel[0]) - return + return s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, responseBody, canonicalModel[0]) } - s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, responseBody) + return s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, responseBody) } func (s *OpenAIGatewayService) handleErrorResponse( @@ -470,7 +469,7 @@ func (s *OpenAIGatewayService) handleErrorResponse( return nil, &UpstreamFailoverError{ StatusCode: resp.StatusCode, ResponseBody: body, - RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), + RetryableOnSameAccount: false, } } @@ -640,7 +639,7 @@ func (s *OpenAIGatewayService) handleCompatErrorResponse( return nil, &UpstreamFailoverError{ StatusCode: resp.StatusCode, ResponseBody: body, - RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), + RetryableOnSameAccount: false, } } diff --git a/backend/internal/service/openai_images.go b/backend/internal/service/openai_images.go index fa0b7f8b6..f32fc5615 100644 --- a/backend/internal/service/openai_images.go +++ b/backend/internal/service/openai_images.go @@ -652,11 +652,11 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesAPIKey( Kind: "failover", Message: upstreamMsg, }) - s.handleFailoverSideEffects(upstreamCtx, resp, account, respBody, upstreamModel) + shouldDisable := s.handleFailoverSideEffects(upstreamCtx, resp, account, respBody, upstreamModel) return nil, &UpstreamFailoverError{ StatusCode: resp.StatusCode, ResponseBody: respBody, - RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), + RetryableOnSameAccount: !shouldDisable && account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), } } return s.handleOpenAIImagesErrorResponse(upstreamCtx, resp, c, account, upstreamModel) diff --git a/backend/internal/service/openai_images_responses.go b/backend/internal/service/openai_images_responses.go index f5b99ac9c..7924de4fc 100644 --- a/backend/internal/service/openai_images_responses.go +++ b/backend/internal/service/openai_images_responses.go @@ -920,7 +920,7 @@ func (s *OpenAIGatewayService) handleOpenAIImagesErrorResponse( return nil, &UpstreamFailoverError{ StatusCode: resp.StatusCode, ResponseBody: body, - RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), + RetryableOnSameAccount: false, } } @@ -1752,11 +1752,11 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesOAuth( Kind: "failover", Message: upstreamMsg, }) - s.handleFailoverSideEffects(upstreamCtx, resp, account, respBody, requestModel) + shouldDisable := s.handleFailoverSideEffects(upstreamCtx, resp, account, respBody, requestModel) return nil, &UpstreamFailoverError{ StatusCode: resp.StatusCode, ResponseBody: respBody, - RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), + RetryableOnSameAccount: !shouldDisable && account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), } } return s.handleOpenAIImagesErrorResponse(upstreamCtx, resp, c, account, requestModel) @@ -1885,11 +1885,11 @@ func (s *OpenAIGatewayService) handleOpenAIImagesOAuthResponseError( } responseBody := openAIImagesUpstreamErrorResponseBody(upstreamErr) - s.handleOpenAIAccountUpstreamError(ctx, account, upstreamErr.StatusCode, headers, responseBody, requestedModel) + shouldDisable := s.handleOpenAIAccountUpstreamError(ctx, account, upstreamErr.StatusCode, headers, responseBody, requestedModel) return &UpstreamFailoverError{ StatusCode: upstreamErr.StatusCode, ResponseBody: responseBody, ResponseHeaders: headers, - RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(upstreamErr.StatusCode), + RetryableOnSameAccount: !shouldDisable && account.IsPoolMode() && account.IsPoolModeRetryableStatus(upstreamErr.StatusCode), } } diff --git a/backend/internal/service/openai_stream_read_error.go b/backend/internal/service/openai_stream_read_error.go new file mode 100644 index 000000000..b318446e0 --- /dev/null +++ b/backend/internal/service/openai_stream_read_error.go @@ -0,0 +1,59 @@ +package service + +import ( + "errors" + "fmt" + "strings" +) + +const ( + // OpenAIUpstreamHTTP2StreamErrorCode is returned to OpenAI-compatible clients + // when an upstream HTTP/2 response stream is reset after the request started. + OpenAIUpstreamHTTP2StreamErrorCode = "upstream_http2_stream_error" + OpenAIUpstreamStreamReadErrorCode = "upstream_stream_read_error" +) + +type openAIUpstreamStreamReadError struct { + cause error + clientCode string + clientMessage string +} + +func (e *openAIUpstreamStreamReadError) Error() string { + return fmt.Sprintf("stream usage incomplete: %v", e.cause) +} + +func (e *openAIUpstreamStreamReadError) Unwrap() error { return e.cause } + +func newOpenAIUpstreamStreamReadError(err error) error { + code, message := classifyOpenAIUpstreamStreamReadError(err) + return &openAIUpstreamStreamReadError{ + cause: err, + clientCode: code, + clientMessage: message, + } +} + +// OpenAIUpstreamStreamReadErrorDetails returns the stable, sanitized client +// classification attached to an upstream stream read failure. +func OpenAIUpstreamStreamReadErrorDetails(err error) (code, message string, ok bool) { + var streamErr *openAIUpstreamStreamReadError + if !errors.As(err, &streamErr) || streamErr == nil { + return "", "", false + } + return streamErr.clientCode, streamErr.clientMessage, true +} + +func classifyOpenAIUpstreamStreamReadError(err error) (code, message string) { + if err != nil { + lower := strings.ToLower(err.Error()) + // net/http's HTTP/2 stream error is unexported. Its stable text contains + // "stream error: stream ID ..."; match only the transport signature and + // never pass the original text to the client. + if strings.Contains(lower, "stream error: stream id ") || + (strings.Contains(lower, "http2:") && strings.Contains(lower, "stream")) { + return OpenAIUpstreamHTTP2StreamErrorCode, "Upstream HTTP/2 stream failed" + } + } + return OpenAIUpstreamStreamReadErrorCode, "Upstream response stream was interrupted" +} diff --git a/backend/internal/service/openai_ws_client_read.go b/backend/internal/service/openai_ws_client_read.go index da8d650be..4c9d07c89 100644 --- a/backend/internal/service/openai_ws_client_read.go +++ b/backend/internal/service/openai_ws_client_read.go @@ -22,6 +22,29 @@ func ReadOpenAIWSClientMessage( timeout time.Duration, timeoutStatus coderws.StatusCode, timeoutReason string, +) (coderws.MessageType, []byte, error) { + return readOpenAIWSClientMessageWithTimeoutStart( + controlCtx, + conn, + timeout, + timeoutStatus, + timeoutReason, + nil, + nil, + ) +} + +// readOpenAIWSClientMessageWithTimeoutStart supports readers whose timeout +// starts after a state transition, such as a completed passthrough turn. When +// timeoutActive is nil, a positive timeout starts immediately. +func readOpenAIWSClientMessageWithTimeoutStart( + controlCtx context.Context, + conn *coderws.Conn, + timeout time.Duration, + timeoutStatus coderws.StatusCode, + timeoutReason string, + timeoutStart <-chan struct{}, + timeoutActive func() bool, ) (coderws.MessageType, []byte, error) { if conn == nil { return 0, nil, errors.New("openai websocket client connection is nil") @@ -36,13 +59,33 @@ func ReadOpenAIWSClientMessage( readDone <- openAIWSClientReadResult{messageType: messageType, payload: payload, err: err} }() - var timeoutCh <-chan time.Time var timer *time.Timer - if timeout > 0 { - timer = time.NewTimer(timeout) + var timeoutCh <-chan time.Time + startTimeout := func() { + if timeout <= 0 || (timeoutActive != nil && !timeoutActive()) { + return + } + if timer == nil { + timer = time.NewTimer(timeout) + } else { + if !timer.Stop() { + select { + case <-timer.C: + default: + } + } + timer.Reset(timeout) + } timeoutCh = timer.C - defer timer.Stop() } + if timeoutActive == nil || timeoutActive() { + startTimeout() + } + defer func() { + if timer != nil { + timer.Stop() + } + }() closeAndJoin := func(status coderws.StatusCode, reason string, cause error) (coderws.MessageType, []byte, error) { _ = conn.Close(status, reason) @@ -51,20 +94,24 @@ func ReadOpenAIWSClientMessage( 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, - ) + for { + select { + case result := <-readDone: + return result.messageType, result.payload, result.err + case <-timeoutStart: + startTimeout() + 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) } - return closeAndJoin(coderws.StatusGoingAway, "websocket request canceled", cause) } } diff --git a/backend/internal/service/openai_ws_v2/passthrough_relay.go b/backend/internal/service/openai_ws_v2/passthrough_relay.go index d41abaac3..4dae079b5 100644 --- a/backend/internal/service/openai_ws_v2/passthrough_relay.go +++ b/backend/internal/service/openai_ws_v2/passthrough_relay.go @@ -52,6 +52,7 @@ type RelayTurnResult struct { type RelayExit struct { Stage string Err error + Graceful bool WroteDownstream bool } @@ -65,6 +66,9 @@ type RelayOptions struct { OnUsageParseFailure func(eventType string, usageRaw string) OnTurnComplete func(turn RelayTurnResult) BeforeWriteClient func(msgType coderws.MessageType, payload []byte, wroteDownstream bool) error + BeforeClientWrite func(msgType coderws.MessageType, payload []byte) + AfterClientWrite func(msgType coderws.MessageType, payload []byte, writeErr error) + BeforeRelayCancel func(exit RelayExit) ReadClientFrame func(ctx context.Context, clientConn FrameConn) (coderws.MessageType, []byte, error) OnTrace func(event RelayTraceEvent) Now func() time.Time @@ -226,7 +230,9 @@ func Relay( options.OnUsageParseFailure, options.OnTurnComplete, options.BeforeWriteClient, - func() { + options.BeforeClientWrite, + options.AfterClientWrite, + func(msgType coderws.MessageType, payload []byte) { if options.StartClientAfterFirstDownstream { startClientReader() } @@ -241,6 +247,13 @@ func Relay( go runIdleWatchdog(relayCtx, nowFn, options.IdleTimeout, &lastActivity, onTrace, exitCh) firstExit := <-exitCh + // An outer ingress cancellation is a control-plane close, not a graceful + // upstream disconnect. Leave the client connection open here so the + // adapter can emit the precise lease/request close code. Internal + // relayCancel does not cancel ctx and therefore does not take this path. + if ctx.Err() != nil { + firstExit.graceful = false + } emitRelayTrace(onTrace, RelayTraceEvent{ Stage: "first_exit", Direction: relayDirectionFromStage(firstExit.stage), @@ -248,6 +261,14 @@ func Relay( WroteDownstream: firstExit.wroteDownstream, Error: relayErrorString(firstExit.err), }) + if options.BeforeRelayCancel != nil { + options.BeforeRelayCancel(RelayExit{ + Stage: firstExit.stage, + Err: firstExit.err, + Graceful: firstExit.graceful, + WroteDownstream: firstExit.wroteDownstream, + }) + } combinedWroteDownstream := firstExit.wroteDownstream secondExit := relayExitSignal{graceful: true} hasSecondExit := false @@ -422,7 +443,9 @@ func runUpstreamToClient( onUsageParseFailure func(eventType string, usageRaw string), onTurnComplete func(turn RelayTurnResult), beforeWriteClient func(msgType coderws.MessageType, payload []byte, wroteDownstream bool) error, - afterWriteClient func(), + beforeClientWrite func(msgType coderws.MessageType, payload []byte), + afterClientWrite func(msgType coderws.MessageType, payload []byte, writeErr error), + afterWriteClient func(msgType coderws.MessageType, payload []byte), dropDownstreamWrites *atomic.Bool, forwardedFrames *atomic.Int64, droppedFrames *atomic.Int64, @@ -498,21 +521,28 @@ func runUpstreamToClient( markActivity() continue } - if err := writeClient(msgType, payload); err != nil { + if beforeClientWrite != nil { + beforeClientWrite(msgType, payload) + } + writeErr := writeClient(msgType, payload) + if afterClientWrite != nil { + afterClientWrite(msgType, payload, writeErr) + } + if writeErr != nil { emitRelayTrace(onTrace, RelayTraceEvent{ Stage: "write_client_failed", Direction: "upstream_to_client", MessageType: relayMessageTypeString(msgType), PayloadBytes: len(payload), WroteDownstream: wroteDownstream, - Error: err.Error(), + Error: writeErr.Error(), }) - exitCh <- relayExitSignal{stage: "write_client", err: err, wroteDownstream: wroteDownstream} + exitCh <- relayExitSignal{stage: "write_client", err: writeErr, wroteDownstream: wroteDownstream} return } wroteDownstream = true if afterWriteClient != nil { - afterWriteClient() + afterWriteClient(msgType, payload) } if forwardedFrames != nil { forwardedFrames.Add(1) diff --git a/backend/internal/service/openai_ws_v2/passthrough_relay_internal_test.go b/backend/internal/service/openai_ws_v2/passthrough_relay_internal_test.go index f92117a2b..6036be745 100644 --- a/backend/internal/service/openai_ws_v2/passthrough_relay_internal_test.go +++ b/backend/internal/service/openai_ws_v2/passthrough_relay_internal_test.go @@ -125,6 +125,8 @@ func TestRunUpstreamToClient_ErrorAndDropPaths(t *testing.T) { nil, nil, nil, + nil, + nil, drop, nil, nil, @@ -156,6 +158,8 @@ func TestRunUpstreamToClient_ErrorAndDropPaths(t *testing.T) { nil, nil, nil, + nil, + nil, drop, nil, nil, @@ -190,6 +194,8 @@ func TestRunUpstreamToClient_ErrorAndDropPaths(t *testing.T) { nil, nil, nil, + nil, + nil, drop, nil, dropped, 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 cdd41a058..a5d77fe0f 100644 --- a/backend/internal/service/openai_ws_v2/passthrough_relay_test.go +++ b/backend/internal/service/openai_ws_v2/passthrough_relay_test.go @@ -31,6 +31,12 @@ type delayedReadFrameConn struct { once sync.Once } +type readStartSpyFrameConn struct { + base FrameConn + started chan struct{} + startOnce sync.Once +} + type closeSpyFrameConn struct { closeCalls atomic.Int32 } @@ -126,6 +132,19 @@ func (c *delayedReadFrameConn) Close() error { return c.base.Close() } +func (c *readStartSpyFrameConn) ReadFrame(ctx context.Context) (coderws.MessageType, []byte, error) { + c.startOnce.Do(func() { close(c.started) }) + return c.base.ReadFrame(ctx) +} + +func (c *readStartSpyFrameConn) WriteFrame(ctx context.Context, msgType coderws.MessageType, payload []byte) error { + return c.base.WriteFrame(ctx, msgType, payload) +} + +func (c *readStartSpyFrameConn) Close() error { + return c.base.Close() +} + func (c *closeSpyFrameConn) ReadFrame(ctx context.Context) (coderws.MessageType, []byte, error) { if ctx == nil { ctx = context.Background() @@ -661,6 +680,51 @@ func TestRelay_ContextCanceled(t *testing.T) { require.NotNil(t, relayExit) } +func TestRelay_DownstreamPreambleStartsClientReader(t *testing.T) { + clientBase := newPassthroughTestFrameConn(nil, false) + clientConn := &readStartSpyFrameConn{base: clientBase, started: make(chan struct{})} + upstreamConn := newPassthroughTestFrameConn(nil, false) + resultCh := make(chan *RelayExit, 1) + + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + go func() { + _, relayExit := Relay( + ctx, + clientConn, + upstreamConn, + []byte(`{"type":"response.create","model":"gpt-5.1"}`), + RelayOptions{ + StartClientAfterFirstDownstream: true, + }, + ) + resultCh <- relayExit + }() + + upstreamConn.readCh <- passthroughTestFrame{ + msgType: coderws.MessageText, + payload: []byte(`{"type":"response.created","response":{"id":"resp_semantic_gate"}}`), + } + require.Eventually(t, func() bool { return len(clientBase.Writes()) == 1 }, time.Second, 10*time.Millisecond) + select { + case <-clientConn.started: + case <-time.After(time.Second): + t.Fatal("response.created did not start the client reader") + } + + upstreamConn.readCh <- passthroughTestFrame{ + msgType: coderws.MessageText, + payload: []byte(`{"type":"response.completed","response":{"id":"resp_semantic_gate","usage":{"input_tokens":1,"output_tokens":1}}}`), + } + _ = upstreamConn.Close() + select { + case relayExit := <-resultCh: + require.Nil(t, relayExit) + case <-time.After(time.Second): + t.Fatal("relay did not finish after terminal event and upstream close") + } +} + func TestRelay_TraceEvents_ContainsLifecycleStages(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 af1bf7456..f61cb572b 100644 --- a/backend/internal/service/openai_ws_v2_passthrough_adapter.go +++ b/backend/internal/service/openai_ws_v2_passthrough_adapter.go @@ -7,7 +7,9 @@ import ( "net/http" "net/url" "strings" + "sync" "sync/atomic" + "time" "github.com/Wei-Shaw/sub2api/internal/pkg/logger" "github.com/Wei-Shaw/sub2api/internal/pkg/openai" @@ -18,7 +20,11 @@ import ( ) type openAIWSClientFrameConn struct { - conn *coderws.Conn + conn *coderws.Conn + controlCtx context.Context + interTurnIdleTimeout time.Duration + interTurnStarted chan struct{} + waitingForNextTurn atomic.Bool } // openAIWSPolicyEnforcingFrameConn wraps a client-side FrameConn and runs @@ -193,16 +199,379 @@ func openAIWSPassthroughRequestModelFromSessionFrame(payload []byte) string { const openaiWSV2PassthroughModeFields = "ws_mode=passthrough ws_router=v2" -var _ openaiwsv2.FrameConn = (*openAIWSClientFrameConn)(nil) +var errOpenAIWSPassthroughFirstOutputTimeout = errors.New("openai websocket passthrough first output timeout") +var errOpenAIWSPassthroughActiveTurnTimeout = errors.New("openai websocket passthrough active turn read timeout") -func (c *openAIWSClientFrameConn) ReadFrame(ctx context.Context) (coderws.MessageType, []byte, error) { - if c == nil || c.conn == nil { +type openAIWSPassthroughDeadlinePhase uint8 + +const ( + openAIWSPassthroughDeadlinePhaseFirstSemantic openAIWSPassthroughDeadlinePhase = iota + 1 + openAIWSPassthroughDeadlinePhaseActiveRead +) + +type openAIWSPassthroughFirstOutputDeadline struct { + timeout time.Duration + startedAt time.Time + requestModel string + reasoningEffort string + phase openAIWSPassthroughDeadlinePhase +} + +type openAIWSPassthroughFirstOutputTimeoutError struct { + deadline openAIWSPassthroughFirstOutputDeadline +} + +func (e *openAIWSPassthroughFirstOutputTimeoutError) Error() string { + return errOpenAIWSPassthroughFirstOutputTimeout.Error() +} + +func (e *openAIWSPassthroughFirstOutputTimeoutError) Unwrap() error { + return errOpenAIWSPassthroughFirstOutputTimeout +} + +type openAIWSPassthroughActiveTurnTimeoutError struct{} + +func (e *openAIWSPassthroughActiveTurnTimeoutError) Error() string { + return errOpenAIWSPassthroughActiveTurnTimeout.Error() +} + +func (e *openAIWSPassthroughActiveTurnTimeoutError) Unwrap() error { + return errOpenAIWSPassthroughActiveTurnTimeout +} + +type openAIWSPassthroughFirstOutputDeadlineState struct { + armed bool + generation uint64 + deadline openAIWSPassthroughFirstOutputDeadline +} + +type openAIWSPassthroughTurnLifecycle struct { + mu sync.Mutex + inFlight bool +} + +func newOpenAIWSPassthroughTurnLifecycle(inFlight bool) *openAIWSPassthroughTurnLifecycle { + return &openAIWSPassthroughTurnLifecycle{inFlight: inFlight} +} + +func (l *openAIWSPassthroughTurnLifecycle) beginResponseCreate(onAccepted func()) bool { + if l == nil { + return false + } + l.mu.Lock() + defer l.mu.Unlock() + if l.inFlight { + return false + } + l.inFlight = true + if onAccepted != nil { + onAccepted() + } + return true +} + +func (l *openAIWSPassthroughTurnLifecycle) cancelResponseCreate() { + if l == nil { + return + } + l.mu.Lock() + l.inFlight = false + l.mu.Unlock() +} + +func (l *openAIWSPassthroughTurnLifecycle) beginTerminalWrite() { + if l != nil { + l.mu.Lock() + } +} + +func (l *openAIWSPassthroughTurnLifecycle) finishTerminalWrite(succeeded bool, onSucceeded func()) { + if l == nil { + return + } + if succeeded { + if onSucceeded != nil { + onSucceeded() + } + l.inFlight = false + } + l.mu.Unlock() +} + +type openAIWSPassthroughFirstOutputFrameConn struct { + inner openaiwsv2.FrameConn + resolveDeadline func(payload []byte) openAIWSPassthroughFirstOutputDeadline + activeReadTimeout time.Duration + + mu sync.Mutex + state openAIWSPassthroughFirstOutputDeadlineState + deadlineChanged chan struct{} +} + +func (c *openAIWSPassthroughFirstOutputFrameConn) ReadFrame(ctx context.Context) (coderws.MessageType, []byte, error) { + if c == nil || c.inner == nil { return coderws.MessageText, nil, errOpenAIWSConnClosed } if ctx == nil { ctx = context.Background() } - return c.conn.Read(ctx) + + type readResult struct { + msgType coderws.MessageType + payload []byte + err error + } + readCtx, cancelRead := context.WithCancel(ctx) + readResultCh := make(chan readResult, 1) + go func() { + msgType, payload, err := c.inner.ReadFrame(readCtx) + readResultCh <- readResult{msgType: msgType, payload: payload, err: err} + }() + + var timer *time.Timer + var timerCh <-chan time.Time + resetTimer := func() { + state := c.deadlineState() + if timer != nil { + if !timer.Stop() { + select { + case <-timer.C: + default: + } + } + } + if !state.armed || state.deadline.timeout <= 0 { + timerCh = nil + return + } + remaining := time.Until(state.deadline.startedAt.Add(state.deadline.timeout)) + if remaining < 0 { + remaining = 0 + } + if timer == nil { + timer = time.NewTimer(remaining) + } else { + timer.Reset(remaining) + } + timerCh = timer.C + } + resetTimer() + + defer func() { + cancelRead() + if timer != nil { + timer.Stop() + } + }() + for { + select { + case result := <-readResultCh: + if result.err == nil { + c.observeUpstreamActivity(result.msgType, result.payload) + } + return result.msgType, result.payload, result.err + case <-c.deadlineChanged: + resetTimer() + case <-timerCh: + state := c.deadlineState() + if !state.armed || state.deadline.timeout <= 0 || time.Now().Before(state.deadline.startedAt.Add(state.deadline.timeout)) { + resetTimer() + continue + } + if ctx.Err() != nil { + cancelRead() + <-readResultCh + return coderws.MessageText, nil, ctx.Err() + } + cancelRead() + <-readResultCh + if state.deadline.phase == openAIWSPassthroughDeadlinePhaseActiveRead { + return coderws.MessageText, nil, &openAIWSPassthroughActiveTurnTimeoutError{} + } + return coderws.MessageText, nil, &openAIWSPassthroughFirstOutputTimeoutError{deadline: state.deadline} + case <-ctx.Done(): + cancelRead() + <-readResultCh + return coderws.MessageText, nil, ctx.Err() + } + } +} + +func (c *openAIWSPassthroughFirstOutputFrameConn) WriteFrame(ctx context.Context, msgType coderws.MessageType, payload []byte) error { + if c == nil || c.inner == nil { + return errOpenAIWSConnClosed + } + generation := uint64(0) + if msgType == coderws.MessageText && strings.TrimSpace(gjson.GetBytes(payload, "type").String()) == "response.create" { + generation = c.armDeadline(payload) + } + if err := c.inner.WriteFrame(ctx, msgType, payload); err != nil { + c.disarmDeadline(generation) + return err + } + return nil +} + +func (c *openAIWSPassthroughFirstOutputFrameConn) Close() error { + if c == nil || c.inner == nil { + return nil + } + return c.inner.Close() +} + +func (c *openAIWSPassthroughFirstOutputFrameConn) armDeadline(payload []byte) uint64 { + if c == nil || c.resolveDeadline == nil { + return 0 + } + deadline := c.resolveDeadline(payload) + if deadline.timeout <= 0 { + return 0 + } + if deadline.startedAt.IsZero() { + deadline.startedAt = time.Now() + } + deadline.phase = openAIWSPassthroughDeadlinePhaseFirstSemantic + c.mu.Lock() + c.state.generation++ + generation := c.state.generation + c.state.armed = true + c.state.deadline = deadline + c.mu.Unlock() + c.notifyDeadlineChanged() + return generation +} + +func (c *openAIWSPassthroughFirstOutputFrameConn) observeUpstreamActivity(msgType coderws.MessageType, payload []byte) { + if c == nil { + return + } + if msgType == coderws.MessageText && openAIWSPassthroughIsTerminalOutput(payload) { + c.disarmDeadline(0) + return + } + state := c.deadlineState() + if state.armed && state.deadline.phase == openAIWSPassthroughDeadlinePhaseActiveRead { + c.armActiveReadDeadline() + return + } + if msgType == coderws.MessageText && openAIWSPassthroughStartsSemanticOutput(payload) { + c.armActiveReadDeadline() + } +} + +func (c *openAIWSPassthroughFirstOutputFrameConn) armActiveReadDeadline() { + if c == nil { + return + } + if c.activeReadTimeout <= 0 { + c.disarmDeadline(0) + return + } + c.mu.Lock() + c.state.generation++ + c.state.armed = true + c.state.deadline = openAIWSPassthroughFirstOutputDeadline{ + timeout: c.activeReadTimeout, + startedAt: time.Now(), + phase: openAIWSPassthroughDeadlinePhaseActiveRead, + } + c.mu.Unlock() + c.notifyDeadlineChanged() +} + +func (c *openAIWSPassthroughFirstOutputFrameConn) disarmDeadline(generation uint64) { + if c == nil { + return + } + c.mu.Lock() + if !c.state.armed || (generation != 0 && generation != c.state.generation) { + c.mu.Unlock() + return + } + c.state.armed = false + c.mu.Unlock() + c.notifyDeadlineChanged() +} + +func (c *openAIWSPassthroughFirstOutputFrameConn) deadlineState() openAIWSPassthroughFirstOutputDeadlineState { + if c == nil { + return openAIWSPassthroughFirstOutputDeadlineState{} + } + c.mu.Lock() + defer c.mu.Unlock() + return c.state +} + +func (c *openAIWSPassthroughFirstOutputFrameConn) notifyDeadlineChanged() { + if c == nil || c.deadlineChanged == nil { + return + } + select { + case c.deadlineChanged <- struct{}{}: + default: + } +} + +func openAIWSPassthroughStartsSemanticOutput(payload []byte) bool { + eventType := strings.TrimSpace(gjson.GetBytes(payload, "type").String()) + switch eventType { + case "response.completed", "response.done", "response.failed", "response.incomplete", "response.cancelled", "response.canceled": + return true + case "", "response.created", "response.in_progress", "response.output_item.added", "response.output_item.done": + return false + } + return strings.Contains(eventType, ".delta") || + strings.HasPrefix(eventType, "response.output_text") || + strings.HasPrefix(eventType, "response.output") +} + +func openAIWSPassthroughIsTerminalOutput(payload []byte) bool { + switch strings.TrimSpace(gjson.GetBytes(payload, "type").String()) { + case "response.completed", "response.done", "response.failed", "response.incomplete", "response.cancelled", "response.canceled": + return true + default: + return false + } +} + +var _ openaiwsv2.FrameConn = (*openAIWSClientFrameConn)(nil) +var _ openaiwsv2.FrameConn = (*openAIWSPassthroughFirstOutputFrameConn)(nil) + +func (c *openAIWSClientFrameConn) ReadFrame(ctx context.Context) (coderws.MessageType, []byte, error) { + if c == nil || c.conn == nil { + return coderws.MessageText, nil, errOpenAIWSConnClosed + } + controlCtx := ctx + if c.controlCtx != nil { + controlCtx = c.controlCtx + } + msgType, payload, err := readOpenAIWSClientMessageWithTimeoutStart( + controlCtx, + c.conn, + c.interTurnIdleTimeout, + coderws.StatusNormalClosure, + "websocket idle timeout", + c.interTurnStarted, + func() bool { return c.waitingForNextTurn.Load() }, + ) + return msgType, payload, err +} + +func (c *openAIWSClientFrameConn) markTurnStarted() { + if c != nil { + c.waitingForNextTurn.Store(false) + } +} + +func (c *openAIWSClientFrameConn) markTurnCompleted() { + if c == nil { + return + } + c.waitingForNextTurn.Store(true) + select { + case c.interTurnStarted <- struct{}{}: + default: + } } func (c *openAIWSClientFrameConn) WriteFrame(ctx context.Context, msgType coderws.MessageType, payload []byte) error { @@ -429,19 +798,68 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough( if !ok { return errors.New("openai ws passthrough upstream connection does not support frame relay") } + relayUpstreamFrameConn := &openAIWSPassthroughFirstOutputFrameConn{ + inner: upstreamFrameConn, + activeReadTimeout: s.openAIWSPassthroughIdleTimeout(), + deadlineChanged: make(chan struct{}, 1), + resolveDeadline: func(payload []byte) openAIWSPassthroughFirstOutputDeadline { + reasoningEffort := "" + if current := usageMeta.reasoningEffort.Load(); current != nil { + reasoningEffort = *current + } + timeout := s.openAIFirstOutputTimeout(reasoningEffort) + if timeout <= 0 { + timeout = s.openAIWSPassthroughIdleTimeout() + } + model := openAIWSPassthroughRequestModelForFrame(payload) + if model == "" { + model = usageMeta.requestModelForFrame(payload) + } + if model == "" { + model = requestModel + } + return openAIWSPassthroughFirstOutputDeadline{ + timeout: timeout, + startedAt: time.Now(), + requestModel: model, + reasoningEffort: reasoningEffort, + } + }, + } completedTurns := atomic.Int32{} + turnLifecycle := newOpenAIWSPassthroughTurnLifecycle(true) + clientFrameConn := &openAIWSClientFrameConn{ + conn: clientConn, + controlCtx: ctx, + interTurnIdleTimeout: s.openAIWSIngressInterTurnIdleTimeout(), + interTurnStarted: make(chan struct{}, 1), + } policyClientConn := &openAIWSPolicyEnforcingFrameConn{ - inner: &openAIWSClientFrameConn{conn: clientConn}, + inner: clientFrameConn, // 注意线程安全:filter 仅在 runClientToUpstream 这一条 // goroutine 中被调用(passthrough_relay.go: ReadFrame loop), // capturedSessionModel 的读写都发生在该 goroutine 内,因此无需 // 加锁/原子化。 - filter: func(msgType coderws.MessageType, payload []byte) ([]byte, *OpenAIFastBlockedError, error) { + filter: func(msgType coderws.MessageType, payload []byte) (out []byte, blocked *OpenAIFastBlockedError, filterErr error) { if msgType != coderws.MessageText { return payload, nil, nil } - if strings.TrimSpace(gjson.GetBytes(payload, "type").String()) == "response.create" { + eventType := strings.TrimSpace(gjson.GetBytes(payload, "type").String()) + isResponseCreate := eventType == "response.create" + acceptedTurn := false + if isResponseCreate { + if !turnLifecycle.beginResponseCreate(clientFrameConn.markTurnStarted) { + err := errors.New("overlapping response.create is not supported") + return payload, nil, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, err.Error(), err) + } + defer func() { + if !acceptedTurn { + turnLifecycle.cancelResponseCreate() + } + }() + } + if isResponseCreate { if account.IsOpenAIOAuth() && isOpenAIResponsesLiteWebSocketPayload(payload) { litePayload, _, liteErr := normalizeOpenAIResponsesLiteToolsPayload(payload) if liteErr != nil { @@ -450,7 +868,7 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough( payload = litePayload } } - if strings.TrimSpace(gjson.GetBytes(payload, "type").String()) == "response.create" && hooks != nil && hooks.BeforeRequest != nil { + if isResponseCreate && hooks != nil && hooks.BeforeRequest != nil { turnNo := int(completedTurns.Load()) + 1 if turnNo < 2 { turnNo = 2 @@ -499,9 +917,9 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough( // extractOpenAIServiceTierFromBody 返回 nil;这里有意 // 覆盖(Store(nil)),因为 OpenAI 上游对该帧实际不传 // service_tier 时按 default 处理,billing 应如实反映。 - if policyErr == nil && blocked == nil && - strings.TrimSpace(gjson.GetBytes(payload, "type").String()) == "response.create" { + if policyErr == nil && blocked == nil && isResponseCreate { usageMeta.updateFromResponseCreate(out, model, requestModelForThisFrame) + acceptedTurn = true } return out, blocked, policyErr }, @@ -521,7 +939,7 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough( } upstreamFirstMessageSent := false firstWriteCtx, cancelFirstWrite := context.WithTimeout(ctx, s.openAIWSWriteTimeout()) - firstWriteErr := upstreamFrameConn.WriteFrame(firstWriteCtx, coderws.MessageText, firstClientMessage) + firstWriteErr := relayUpstreamFrameConn.WriteFrame(firstWriteCtx, coderws.MessageText, firstClientMessage) cancelFirstWrite() if firstWriteErr != nil { return wrapOpenAIWSIngressTurnError( @@ -550,11 +968,14 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough( relayResult, relayExit := openaiwsv2.RunEntry(openaiwsv2.EntryInput{ Ctx: ctx, ClientConn: policyClientConn, - UpstreamConn: upstreamFrameConn, + UpstreamConn: relayUpstreamFrameConn, FirstClientMessage: firstClientMessage, Options: openaiwsv2.RelayOptions{ - WriteTimeout: s.openAIWSWriteTimeout(), - IdleTimeout: s.openAIWSPassthroughIdleTimeout(), + WriteTimeout: s.openAIWSWriteTimeout(), + // Passthrough idle is enforced only after a completed turn by + // clientFrameConn. The relay-wide activity watchdog would also + // terminate a healthy active upstream turn. + IdleTimeout: 0, FirstMessageType: coderws.MessageText, FirstMessageSent: upstreamFirstMessageSent, StartClientAfterFirstDownstream: true, @@ -603,6 +1024,27 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough( hooks.AfterTurn(turnNo, turnResult, nil) } }, + BeforeClientWrite: func(msgType coderws.MessageType, payload []byte) { + if msgType == coderws.MessageText && openAIWSPassthroughIsTerminalOutput(payload) { + turnLifecycle.beginTerminalWrite() + } + }, + AfterClientWrite: func(msgType coderws.MessageType, payload []byte, writeErr error) { + if msgType == coderws.MessageText && openAIWSPassthroughIsTerminalOutput(payload) { + turnLifecycle.finishTerminalWrite(writeErr == nil, clientFrameConn.markTurnCompleted) + } + }, + BeforeRelayCancel: func(exit openaiwsv2.RelayExit) { + if context.Cause(ctx) != nil { + return + } + status, reason, ok := openAIWSPassthroughRelayClientClose(exit, int(completedTurns.Load())) + if !ok { + return + } + _ = clientConn.Close(status, reason) + _ = clientConn.CloseNow() + }, BeforeWriteClient: func(msgType coderws.MessageType, payload []byte, wroteDownstream bool) error { if msgType != coderws.MessageText { return nil @@ -650,6 +1092,17 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough( }, }, }) + if cause := context.Cause(ctx); cause != nil { + status := coderws.StatusGoingAway + reason := "websocket request canceled" + if errors.Is(cause, ErrOpenAIWSIngressLeaseLost) { + status = coderws.StatusTryAgainLater + reason = "websocket ingress capacity lease lost; please reconnect" + } + _ = clientConn.Close(status, reason) + _ = clientConn.CloseNow() + return NewOpenAIWSClientCloseError(status, reason, cause) + } result := &OpenAIForwardResult{ RequestID: relayResult.RequestID, @@ -704,6 +1157,41 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough( ) relayErr := relayExit.Err + var firstOutputTimeoutErr *openAIWSPassthroughFirstOutputTimeoutError + if errors.As(relayErr, &firstOutputTimeoutErr) { + deadline := firstOutputTimeoutErr.deadline + failoverErr := s.newOpenAIFirstOutputTimeoutError( + ctx, + c, + account, + deadline.startedAt, + deadline.requestModel, + deadline.reasoningEffort, + deadline.timeout, + "websocket_first_semantic_output", + handshakeHeaders, + ) + if turnCount == 0 && !relayExit.WroteDownstream { + relayErr = failoverErr + } else { + // The handler only retains the initial response.create across + // account attempts. Replaying it after a later-turn timeout would + // duplicate the first turn, so later turns end the client session. + relayErr = NewOpenAIWSClientCloseError( + coderws.StatusGoingAway, + "upstream produced no semantic output; please reconnect", + firstOutputTimeoutErr, + ) + } + } + var activeTurnTimeoutErr *openAIWSPassthroughActiveTurnTimeoutError + if errors.As(relayErr, &activeTurnTimeoutErr) { + relayErr = NewOpenAIWSClientCloseError( + coderws.StatusGoingAway, + "upstream websocket read timeout; please reconnect", + activeTurnTimeoutErr, + ) + } if relayExit.Stage == "idle_timeout" { relayErr = NewOpenAIWSClientCloseError( coderws.StatusPolicyViolation, @@ -722,6 +1210,28 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough( return turnErr } +func openAIWSPassthroughRelayClientClose(exit openaiwsv2.RelayExit, completedTurns int) (coderws.StatusCode, string, bool) { + var closeErr *OpenAIWSClientCloseError + if errors.As(exit.Err, &closeErr) { + return closeErr.StatusCode(), closeErr.Reason(), true + } + var activeTurnTimeoutErr *openAIWSPassthroughActiveTurnTimeoutError + if errors.As(exit.Err, &activeTurnTimeoutErr) { + return coderws.StatusGoingAway, "upstream websocket read timeout; please reconnect", true + } + var firstOutputTimeoutErr *openAIWSPassthroughFirstOutputTimeoutError + if errors.As(exit.Err, &firstOutputTimeoutErr) { + if completedTurns > 0 || exit.WroteDownstream { + return coderws.StatusGoingAway, "upstream produced no semantic output; please reconnect", true + } + return 0, "", false + } + if !exit.Graceful && exit.Stage == "read_upstream" { + return coderws.StatusInternalError, "upstream websocket proxy failed", true + } + return 0, "", false +} + func (s *OpenAIGatewayService) mapOpenAIWSPassthroughDialError( err error, statusCode int, diff --git a/backend/internal/service/openai_ws_v2_passthrough_lifecycle_test.go b/backend/internal/service/openai_ws_v2_passthrough_lifecycle_test.go new file mode 100644 index 000000000..de3d5fdf2 --- /dev/null +++ b/backend/internal/service/openai_ws_v2_passthrough_lifecycle_test.go @@ -0,0 +1,591 @@ +package service + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + coderws "github.com/coder/websocket" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +type stagedPassthroughFrame struct { + messageType coderws.MessageType + payload []byte +} + +type stagedPassthroughConn struct { + frames chan stagedPassthroughFrame + writes chan []byte + closed chan struct{} + closeOnce sync.Once +} + +func newStagedPassthroughConn() *stagedPassthroughConn { + return &stagedPassthroughConn{ + frames: make(chan stagedPassthroughFrame, 4), + writes: make(chan []byte, 4), + closed: make(chan struct{}), + } +} + +func (c *stagedPassthroughConn) Send(payload string) { + c.frames <- stagedPassthroughFrame{messageType: coderws.MessageText, payload: []byte(payload)} +} + +func (c *stagedPassthroughConn) WriteJSON(context.Context, any) error { return nil } + +func (c *stagedPassthroughConn) ReadMessage(ctx context.Context) ([]byte, error) { + _, payload, err := c.ReadFrame(ctx) + return payload, err +} + +func (c *stagedPassthroughConn) Ping(context.Context) error { return nil } + +func (c *stagedPassthroughConn) ReadFrame(ctx context.Context) (coderws.MessageType, []byte, error) { + if ctx == nil { + ctx = context.Background() + } + select { + case <-ctx.Done(): + return coderws.MessageText, nil, ctx.Err() + case <-c.closed: + return coderws.MessageText, nil, errOpenAIWSConnClosed + case frame := <-c.frames: + return frame.messageType, append([]byte(nil), frame.payload...), nil + } +} + +func (c *stagedPassthroughConn) WriteFrame(ctx context.Context, _ coderws.MessageType, payload []byte) error { + if ctx == nil { + ctx = context.Background() + } + select { + case <-ctx.Done(): + return ctx.Err() + case <-c.closed: + return errOpenAIWSConnClosed + default: + } + var parsed any + if err := json.Unmarshal(payload, &parsed); err != nil { + return err + } + select { + case c.writes <- append([]byte(nil), payload...): + case <-ctx.Done(): + return ctx.Err() + case <-c.closed: + return errOpenAIWSConnClosed + } + return nil +} + +func (c *stagedPassthroughConn) Close() error { + c.closeOnce.Do(func() { close(c.closed) }) + return nil +} + +type stagedPassthroughDialer struct { + conn openAIWSClientConn +} + +func (d *stagedPassthroughDialer) Dial(context.Context, string, http.Header, string) (openAIWSClientConn, int, http.Header, error) { + return d.conn, http.StatusSwitchingProtocols, http.Header{}, nil +} + +func newPassthroughLifecycleService(cfg *config.Config, upstream *stagedPassthroughConn) *OpenAIGatewayService { + return &OpenAIGatewayService{ + cfg: cfg, + httpUpstream: &httpUpstreamRecorder{}, + cache: &stubGatewayCache{}, + openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), + toolCorrector: NewCodexToolCorrector(), + openaiWSPassthroughDialer: &stagedPassthroughDialer{conn: upstream}, + } +} + +func passthroughLifecycleConfig() *config.Config { + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + cfg.Gateway.OpenAIFirstOutputTimeoutSeconds = 1 + cfg.Gateway.OpenAIWS.Enabled = true + cfg.Gateway.OpenAIWS.APIKeyEnabled = true + cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true + cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true + cfg.Gateway.OpenAIWS.IngressModeDefault = OpenAIWSIngressModeCtxPool + cfg.Gateway.OpenAIWS.IngressInterTurnIdleTimeoutSeconds = 1 + cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 1 + cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3 + return cfg +} + +func passthroughLifecycleAccount() *Account { + return &Account{ + ID: 901, + Name: "passthrough-lifecycle", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Credentials: map[string]any{"api_key": "sk-test"}, + Extra: map[string]any{ + "openai_apikey_responses_websockets_v2_mode": OpenAIWSIngressModePassthrough, + }, + } +} + +func startPassthroughLifecycleServer( + t *testing.T, + controlCtx context.Context, + svc *OpenAIGatewayService, + account *Account, +) (*httptest.Server, <-chan error) { + t.Helper() + serverErr := make(chan error, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover}) + if err != nil { + serverErr <- err + return + } + defer func() { _ = conn.CloseNow() }() + + msgType, firstMessage, err := ReadOpenAIWSClientMessage( + controlCtx, + conn, + 3*time.Second, + coderws.StatusPolicyViolation, + "missing first response.create message", + ) + if err != nil { + serverErr <- err + return + } + if msgType != coderws.MessageText { + serverErr <- errors.New("first message was not text") + return + } + + recorder := httptest.NewRecorder() + ginCtx, _ := gin.CreateTestContext(recorder) + req := r.Clone(controlCtx) + req.Header = req.Header.Clone() + ginCtx.Request = req + serverErr <- svc.ProxyResponsesWebSocketFromClient(controlCtx, ginCtx, conn, account, "sk-test", firstMessage, nil) + })) + return server, serverErr +} + +func dialPassthroughLifecycleClient(t *testing.T, server *httptest.Server) *coderws.Conn { + t.Helper() + dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second) + clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(server.URL, "http"), nil) + cancelDial() + require.NoError(t, err) + writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second) + err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1","stream":false}`)) + cancelWrite() + require.NoError(t, err) + return clientConn +} + +func readPassthroughLifecycleFrame(t *testing.T, clientConn *coderws.Conn, timeout time.Duration) ([]byte, error) { + t.Helper() + readCtx, cancelRead := context.WithTimeout(context.Background(), timeout) + _, payload, err := clientConn.Read(readCtx) + cancelRead() + return payload, err +} + +func requirePassthroughUpstreamWrite(t *testing.T, upstream *stagedPassthroughConn, timeout time.Duration) []byte { + t.Helper() + select { + case payload := <-upstream.writes: + return payload + case <-time.After(timeout): + t.Fatal("passthrough request was not forwarded upstream") + return nil + } +} + +func TestOpenAIWSPassthroughTurnLifecycle_SerializesTerminalCommitAndNextTurn(t *testing.T) { + clientFrameConn := &openAIWSClientFrameConn{interTurnStarted: make(chan struct{}, 1)} + clientFrameConn.markTurnCompleted() + lifecycle := newOpenAIWSPassthroughTurnLifecycle(true) + lifecycle.beginTerminalWrite() + + admitted := make(chan bool, 1) + go func() { + admitted <- lifecycle.beginResponseCreate(clientFrameConn.markTurnStarted) + }() + select { + case <-admitted: + t.Fatal("next response.create was admitted before terminal commit completed") + case <-time.After(50 * time.Millisecond): + } + + lifecycle.finishTerminalWrite(true, clientFrameConn.markTurnCompleted) + select { + case ok := <-admitted: + require.True(t, ok) + case <-time.After(time.Second): + t.Fatal("next response.create remained blocked after terminal commit") + } + require.False(t, clientFrameConn.waitingForNextTurn.Load(), "accepted next turn must win over terminal idle state") + + lifecycle = newOpenAIWSPassthroughTurnLifecycle(true) + lifecycle.beginTerminalWrite() + admitted = make(chan bool, 1) + go func() { + admitted <- lifecycle.beginResponseCreate(nil) + }() + lifecycle.finishTerminalWrite(false, func() { + t.Error("failed terminal write must not commit idle state") + }) + require.False(t, <-admitted, "failed terminal write must keep the current turn in flight") +} + +func TestPassthroughLifecycle_LeaseLossSendsRetryClose(t *testing.T) { + gin.SetMode(gin.TestMode) + controlCtx, cancelControl := context.WithCancelCause(context.Background()) + upstream := newStagedPassthroughConn() + upstream.Send(`{"type":"response.created","response":{"id":"resp_lease","model":"gpt-5.1"}}`) + server, serverErr := startPassthroughLifecycleServer(t, controlCtx, newPassthroughLifecycleService(passthroughLifecycleConfig(), upstream), passthroughLifecycleAccount()) + defer server.Close() + clientConn := dialPassthroughLifecycleClient(t, server) + defer func() { _ = clientConn.CloseNow() }() + + event, err := readPassthroughLifecycleFrame(t, clientConn, 3*time.Second) + require.NoError(t, err) + require.Equal(t, "response.created", gjson.GetBytes(event, "type").String()) + cancelControl(ErrOpenAIWSIngressLeaseLost) + + _, err = readPassthroughLifecycleFrame(t, clientConn, 3*time.Second) + var closeErr coderws.CloseError + require.ErrorAs(t, err, &closeErr) + require.Equal(t, coderws.StatusTryAgainLater, closeErr.Code) + require.Equal(t, "websocket ingress capacity lease lost; please reconnect", closeErr.Reason) + select { + case <-serverErr: + case <-time.After(3 * time.Second): + t.Fatal("passthrough lease-loss reader did not exit") + } +} + +func TestPassthroughLifecycle_CompletedTurnStartsInterTurnIdle(t *testing.T) { + gin.SetMode(gin.TestMode) + controlCtx, cancelControl := context.WithCancelCause(context.Background()) + defer cancelControl(context.Canceled) + upstream := newStagedPassthroughConn() + upstream.Send(`{"type":"response.completed","response":{"id":"resp_idle","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`) + server, serverErr := startPassthroughLifecycleServer(t, controlCtx, newPassthroughLifecycleService(passthroughLifecycleConfig(), upstream), passthroughLifecycleAccount()) + defer server.Close() + clientConn := dialPassthroughLifecycleClient(t, server) + defer func() { _ = clientConn.CloseNow() }() + + event, err := readPassthroughLifecycleFrame(t, clientConn, 3*time.Second) + require.NoError(t, err) + require.Equal(t, "response.completed", gjson.GetBytes(event, "type").String()) + _, err = readPassthroughLifecycleFrame(t, clientConn, 3*time.Second) + var closeErr coderws.CloseError + require.ErrorAs(t, err, &closeErr) + require.Equal(t, coderws.StatusNormalClosure, closeErr.Code) + require.Equal(t, "websocket idle timeout", closeErr.Reason) + select { + case <-serverErr: + case <-time.After(3 * time.Second): + t.Fatal("passthrough idle reader did not exit") + } +} + +func TestPassthroughLifecycle_ActiveTurnInactivityUsesReadTimeout(t *testing.T) { + gin.SetMode(gin.TestMode) + controlCtx, cancelControl := context.WithCancelCause(context.Background()) + defer cancelControl(context.Canceled) + upstream := newStagedPassthroughConn() + upstream.Send(`{"type":"response.output_text.delta","response_id":"resp_active","delta":"hello"}`) + server, serverErr := startPassthroughLifecycleServer(t, controlCtx, newPassthroughLifecycleService(passthroughLifecycleConfig(), upstream), passthroughLifecycleAccount()) + defer server.Close() + clientConn := dialPassthroughLifecycleClient(t, server) + defer func() { _ = clientConn.CloseNow() }() + + delta, err := readPassthroughLifecycleFrame(t, clientConn, 3*time.Second) + require.NoError(t, err) + require.Equal(t, "response.output_text.delta", gjson.GetBytes(delta, "type").String()) + _, err = readPassthroughLifecycleFrame(t, clientConn, 2500*time.Millisecond) + var websocketCloseErr coderws.CloseError + require.ErrorAs(t, err, &websocketCloseErr) + require.Equal(t, coderws.StatusGoingAway, websocketCloseErr.Code) + require.Equal(t, "upstream websocket read timeout; please reconnect", websocketCloseErr.Reason) + select { + case err := <-serverErr: + var closeErr *OpenAIWSClientCloseError + require.ErrorAs(t, err, &closeErr) + require.Equal(t, coderws.StatusGoingAway, closeErr.StatusCode()) + require.Equal(t, "upstream websocket read timeout; please reconnect", closeErr.Reason()) + case <-time.After(2500 * time.Millisecond): + t.Fatal("passthrough active turn remained unbounded after upstream activity stopped") + } +} + +func TestPassthroughLifecycle_PreambleAllowsPromptClientCancel(t *testing.T) { + gin.SetMode(gin.TestMode) + controlCtx, cancelControl := context.WithCancelCause(context.Background()) + defer cancelControl(context.Canceled) + cfg := passthroughLifecycleConfig() + cfg.Gateway.OpenAIFirstOutputTimeoutSeconds = 3 + upstream := newStagedPassthroughConn() + upstream.Send(`{"type":"response.created","response":{"id":"resp_cancel","model":"gpt-5.1"}}`) + server, serverErr := startPassthroughLifecycleServer(t, controlCtx, newPassthroughLifecycleService(cfg, upstream), passthroughLifecycleAccount()) + defer server.Close() + clientConn := dialPassthroughLifecycleClient(t, server) + defer func() { _ = clientConn.CloseNow() }() + require.Equal(t, "response.create", gjson.GetBytes(requirePassthroughUpstreamWrite(t, upstream, time.Second), "type").String()) + + created, err := readPassthroughLifecycleFrame(t, clientConn, time.Second) + require.NoError(t, err) + require.Equal(t, "response.created", gjson.GetBytes(created, "type").String()) + writeCtx, cancelWrite := context.WithTimeout(context.Background(), time.Second) + err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.cancel","response_id":"resp_cancel"}`)) + cancelWrite() + require.NoError(t, err) + cancelFrame := requirePassthroughUpstreamWrite(t, upstream, 500*time.Millisecond) + require.Equal(t, "response.cancel", gjson.GetBytes(cancelFrame, "type").String()) + + require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done")) + select { + case <-serverErr: + case <-time.After(3 * time.Second): + t.Fatal("passthrough cancel test did not exit") + } +} + +func TestPassthroughLifecycle_RejectsOverlappingResponseCreate(t *testing.T) { + gin.SetMode(gin.TestMode) + controlCtx, cancelControl := context.WithCancelCause(context.Background()) + defer cancelControl(context.Canceled) + cfg := passthroughLifecycleConfig() + cfg.Gateway.OpenAIFirstOutputTimeoutSeconds = 3 + upstream := newStagedPassthroughConn() + upstream.Send(`{"type":"response.created","response":{"id":"resp_overlap_first","model":"gpt-5.1"}}`) + server, serverErr := startPassthroughLifecycleServer(t, controlCtx, newPassthroughLifecycleService(cfg, upstream), passthroughLifecycleAccount()) + defer server.Close() + clientConn := dialPassthroughLifecycleClient(t, server) + defer func() { _ = clientConn.CloseNow() }() + require.Equal(t, "response.create", gjson.GetBytes(requirePassthroughUpstreamWrite(t, upstream, time.Second), "type").String()) + + created, err := readPassthroughLifecycleFrame(t, clientConn, time.Second) + require.NoError(t, err) + require.Equal(t, "response.created", gjson.GetBytes(created, "type").String()) + writeCtx, cancelWrite := context.WithTimeout(context.Background(), time.Second) + err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1"}`)) + cancelWrite() + require.NoError(t, err) + + _, err = readPassthroughLifecycleFrame(t, clientConn, time.Second) + var websocketCloseErr coderws.CloseError + require.ErrorAs(t, err, &websocketCloseErr) + require.Equal(t, coderws.StatusPolicyViolation, websocketCloseErr.Code) + require.Equal(t, "overlapping response.create is not supported", websocketCloseErr.Reason) + select { + case err := <-serverErr: + var closeErr *OpenAIWSClientCloseError + require.ErrorAs(t, err, &closeErr) + require.Equal(t, coderws.StatusPolicyViolation, closeErr.StatusCode()) + require.Equal(t, "overlapping response.create is not supported", closeErr.Reason()) + case <-time.After(3 * time.Second): + t.Fatal("overlapping response.create did not terminate passthrough") + } +} + +func TestPassthroughLifecycle_ActiveTurnActivityRefreshesReadTimeout(t *testing.T) { + gin.SetMode(gin.TestMode) + controlCtx, cancelControl := context.WithCancelCause(context.Background()) + defer cancelControl(context.Canceled) + upstream := newStagedPassthroughConn() + upstream.Send(`{"type":"response.output_text.delta","response_id":"resp_active_refresh","delta":"one"}`) + go func() { + for _, event := range []string{ + `{"type":"response.output_text.delta","response_id":"resp_active_refresh","delta":"two"}`, + `{"type":"response.output_text.delta","response_id":"resp_active_refresh","delta":"three"}`, + `{"type":"response.completed","response":{"id":"resp_active_refresh","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":3}}}`, + } { + timer := time.NewTimer(600 * time.Millisecond) + <-timer.C + timer.Stop() + upstream.Send(event) + } + }() + server, serverErr := startPassthroughLifecycleServer(t, controlCtx, newPassthroughLifecycleService(passthroughLifecycleConfig(), upstream), passthroughLifecycleAccount()) + defer server.Close() + clientConn := dialPassthroughLifecycleClient(t, server) + defer func() { _ = clientConn.CloseNow() }() + + for _, wantType := range []string{ + "response.output_text.delta", + "response.output_text.delta", + "response.output_text.delta", + "response.completed", + } { + frame, err := readPassthroughLifecycleFrame(t, clientConn, 3*time.Second) + require.NoError(t, err) + require.Equal(t, wantType, gjson.GetBytes(frame, "type").String()) + } + require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done")) + select { + case <-serverErr: + case <-time.After(3 * time.Second): + t.Fatal("passthrough active-turn refresh test did not exit") + } +} + +func TestPassthroughLifecycle_TerminalSwitchesToInterTurnIdleTimeout(t *testing.T) { + gin.SetMode(gin.TestMode) + controlCtx, cancelControl := context.WithCancelCause(context.Background()) + defer cancelControl(context.Canceled) + cfg := passthroughLifecycleConfig() + cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 1 + cfg.Gateway.OpenAIWS.IngressInterTurnIdleTimeoutSeconds = 2 + upstream := newStagedPassthroughConn() + upstream.Send(`{"type":"response.completed","response":{"id":"resp_idle_first","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`) + server, serverErr := startPassthroughLifecycleServer(t, controlCtx, newPassthroughLifecycleService(cfg, upstream), passthroughLifecycleAccount()) + defer server.Close() + clientConn := dialPassthroughLifecycleClient(t, server) + defer func() { _ = clientConn.CloseNow() }() + require.Equal(t, "response.create", gjson.GetBytes(requirePassthroughUpstreamWrite(t, upstream, 3*time.Second), "type").String()) + + completed, err := readPassthroughLifecycleFrame(t, clientConn, 3*time.Second) + require.NoError(t, err) + require.Equal(t, "resp_idle_first", gjson.GetBytes(completed, "response.id").String()) + time.Sleep(1300 * time.Millisecond) + writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second) + err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1","previous_response_id":"resp_idle_first"}`)) + cancelWrite() + require.NoError(t, err) + require.Equal(t, "response.create", gjson.GetBytes(requirePassthroughUpstreamWrite(t, upstream, 3*time.Second), "type").String()) + upstream.Send(`{"type":"response.completed","response":{"id":"resp_idle_second","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`) + completed, err = readPassthroughLifecycleFrame(t, clientConn, 3*time.Second) + require.NoError(t, err) + require.Equal(t, "resp_idle_second", gjson.GetBytes(completed, "response.id").String()) + _, err = readPassthroughLifecycleFrame(t, clientConn, 3*time.Second) + var websocketCloseErr coderws.CloseError + require.ErrorAs(t, err, &websocketCloseErr) + require.Equal(t, coderws.StatusNormalClosure, websocketCloseErr.Code) + require.Equal(t, "websocket idle timeout", websocketCloseErr.Reason) + + select { + case err := <-serverErr: + var closeErr *OpenAIWSClientCloseError + require.ErrorAs(t, err, &closeErr) + require.Equal(t, coderws.StatusNormalClosure, closeErr.StatusCode()) + require.Equal(t, "websocket idle timeout", closeErr.Reason()) + case <-time.After(3 * time.Second): + t.Fatal("passthrough terminal turn did not use inter-turn idle timeout") + } +} + +func TestPassthroughLifecycle_FirstOutputTimeoutRemainsBounded(t *testing.T) { + gin.SetMode(gin.TestMode) + controlCtx, cancelControl := context.WithCancelCause(context.Background()) + defer cancelControl(context.Canceled) + upstream := newStagedPassthroughConn() + server, serverErr := startPassthroughLifecycleServer(t, controlCtx, newPassthroughLifecycleService(passthroughLifecycleConfig(), upstream), passthroughLifecycleAccount()) + defer server.Close() + clientConn := dialPassthroughLifecycleClient(t, server) + defer func() { _ = clientConn.CloseNow() }() + + select { + case err := <-serverErr: + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, http.StatusGatewayTimeout, failoverErr.StatusCode) + require.Contains(t, string(failoverErr.ResponseBody), "first_output_timeout") + case <-time.After(2500 * time.Millisecond): + t.Fatal("passthrough first output was left unbounded") + } +} + +func TestPassthroughLifecycle_ResponseCreatedTimeoutClosesWithoutFailover(t *testing.T) { + gin.SetMode(gin.TestMode) + controlCtx, cancelControl := context.WithCancelCause(context.Background()) + defer cancelControl(context.Canceled) + upstream := newStagedPassthroughConn() + upstream.Send(`{"type":"response.created","response":{"id":"resp_preamble","model":"gpt-5.1"}}`) + server, serverErr := startPassthroughLifecycleServer(t, controlCtx, newPassthroughLifecycleService(passthroughLifecycleConfig(), upstream), passthroughLifecycleAccount()) + defer server.Close() + clientConn := dialPassthroughLifecycleClient(t, server) + defer func() { _ = clientConn.CloseNow() }() + + created, err := readPassthroughLifecycleFrame(t, clientConn, 3*time.Second) + require.NoError(t, err) + require.Equal(t, "response.created", gjson.GetBytes(created, "type").String()) + _, err = readPassthroughLifecycleFrame(t, clientConn, 2500*time.Millisecond) + var websocketCloseErr coderws.CloseError + require.ErrorAs(t, err, &websocketCloseErr) + require.Equal(t, coderws.StatusGoingAway, websocketCloseErr.Code) + require.Equal(t, "upstream produced no semantic output; please reconnect", websocketCloseErr.Reason) + select { + case err := <-serverErr: + var failoverErr *UpstreamFailoverError + require.NotErrorAs(t, err, &failoverErr) + var closeErr *OpenAIWSClientCloseError + require.ErrorAs(t, err, &closeErr) + require.Equal(t, coderws.StatusGoingAway, closeErr.StatusCode()) + require.Equal(t, "upstream produced no semantic output; please reconnect", closeErr.Reason()) + case <-time.After(2500 * time.Millisecond): + t.Fatal("response.created timeout did not close the passthrough connection") + } +} + +func TestPassthroughLifecycle_SecondTurnTimeoutIsNotFailoverSafe(t *testing.T) { + gin.SetMode(gin.TestMode) + controlCtx, cancelControl := context.WithCancelCause(context.Background()) + defer cancelControl(context.Canceled) + upstream := newStagedPassthroughConn() + upstream.Send(`{"type":"response.completed","response":{"id":"resp_first","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`) + server, serverErr := startPassthroughLifecycleServer(t, controlCtx, newPassthroughLifecycleService(passthroughLifecycleConfig(), upstream), passthroughLifecycleAccount()) + defer server.Close() + clientConn := dialPassthroughLifecycleClient(t, server) + defer func() { _ = clientConn.CloseNow() }() + + completed, err := readPassthroughLifecycleFrame(t, clientConn, 3*time.Second) + require.NoError(t, err) + require.Equal(t, "response.completed", gjson.GetBytes(completed, "type").String()) + writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second) + err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1","previous_response_id":"resp_first"}`)) + cancelWrite() + require.NoError(t, err) + upstream.Send(`{"type":"response.created","response":{"id":"resp_second","model":"gpt-5.1"}}`) + + created, err := readPassthroughLifecycleFrame(t, clientConn, 3*time.Second) + require.NoError(t, err) + require.Equal(t, "response.created", gjson.GetBytes(created, "type").String()) + _, err = readPassthroughLifecycleFrame(t, clientConn, 2500*time.Millisecond) + var websocketCloseErr coderws.CloseError + require.ErrorAs(t, err, &websocketCloseErr) + require.Equal(t, coderws.StatusGoingAway, websocketCloseErr.Code) + require.Equal(t, "upstream produced no semantic output; please reconnect", websocketCloseErr.Reason) + select { + case err := <-serverErr: + var failoverErr *UpstreamFailoverError + require.NotErrorAs(t, err, &failoverErr, "handler must not replay the initial request on another account for a later-turn timeout") + var closeErr *OpenAIWSClientCloseError + require.ErrorAs(t, err, &closeErr) + require.Equal(t, coderws.StatusGoingAway, closeErr.StatusCode()) + case <-time.After(2500 * time.Millisecond): + t.Fatal("second turn first semantic output was left unbounded") + } +} diff --git a/backend/internal/service/ops_cleanup_executor.go b/backend/internal/service/ops_cleanup_executor.go index d51863c48..a5da61e1d 100644 --- a/backend/internal/service/ops_cleanup_executor.go +++ b/backend/internal/service/ops_cleanup_executor.go @@ -25,19 +25,21 @@ type opsCleanupTarget struct { } type opsCleanupDeletedCounts struct { - errorLogs int64 - alertEvents int64 - systemLogs int64 - logAudits int64 - systemMetrics int64 - hourlyPreagg int64 - dailyPreagg int64 + errorLogs int64 + ingressRejects int64 + alertEvents int64 + systemLogs int64 + logAudits int64 + systemMetrics int64 + hourlyPreagg int64 + dailyPreagg int64 } func (c opsCleanupDeletedCounts) String() string { return fmt.Sprintf( - "error_logs=%d alert_events=%d system_logs=%d log_audits=%d system_metrics=%d hourly_preagg=%d daily_preagg=%d", + "error_logs=%d ingress_rejects=%d alert_events=%d system_logs=%d log_audits=%d system_metrics=%d hourly_preagg=%d daily_preagg=%d", c.errorLogs, + c.ingressRejects, c.alertEvents, c.systemLogs, c.logAudits, diff --git a/backend/internal/service/ops_cleanup_service.go b/backend/internal/service/ops_cleanup_service.go index f812c2906..cfd2f1eed 100644 --- a/backend/internal/service/ops_cleanup_service.go +++ b/backend/internal/service/ops_cleanup_service.go @@ -299,6 +299,7 @@ func (s *OpsCleanupService) runCleanupOnce(ctx context.Context) (opsCleanupDelet targets := []opsCleanupTarget{ {effective.ErrorLogRetentionDays, "ops_error_logs", "created_at", false, &out.errorLogs}, + {effective.ErrorLogRetentionDays, "ops_ingress_reject_aggregates", "bucket_start", false, &out.ingressRejects}, {effective.ErrorLogRetentionDays, "ops_alert_events", "created_at", false, &out.alertEvents}, {effective.ErrorLogRetentionDays, "ops_system_logs", "created_at", false, &out.systemLogs}, {effective.ErrorLogRetentionDays, "ops_system_log_cleanup_audits", "created_at", false, &out.logAudits}, diff --git a/backend/internal/service/ops_ingress_reject.go b/backend/internal/service/ops_ingress_reject.go new file mode 100644 index 000000000..7e0452e5e --- /dev/null +++ b/backend/internal/service/ops_ingress_reject.go @@ -0,0 +1,451 @@ +package service + +import ( + "context" + "log" + "strings" + "sync" + "sync/atomic" + "time" +) + +const ( + ingressRejectShardCount = 16 + ingressRejectMaxEntries = 8192 + ingressRejectMaxPendingBatches = 4 + ingressRejectBucketSize = time.Minute + ingressRejectFlushInterval = 5 * time.Second + ingressRejectFlushTimeout = 5 * time.Second +) + +type OpsIngressRejectAggregate struct { + ID int64 `json:"id"` + BucketStart time.Time `json:"bucket_start"` + RejectReason string `json:"reject_reason"` + RouteFamily string `json:"route_family"` + Protocol string `json:"protocol"` + ClientIP string `json:"client_ip"` + UserID *int64 `json:"user_id,omitempty"` + APIKeyID *int64 `json:"api_key_id,omitempty"` + RequestCount int64 `json:"request_count"` + FirstSeen time.Time `json:"first_seen"` + LastSeen time.Time `json:"last_seen"` +} + +type OpsIngressRejectFilter struct { + StartTime *time.Time + EndTime *time.Time + RejectReason string + RouteFamily string + Protocol string + ClientIP string + UserID *int64 + APIKeyID *int64 + Page int + PageSize int +} + +type OpsIngressRejectList struct { + Items []*OpsIngressRejectAggregate `json:"items"` + Total int `json:"total"` + Page int `json:"page"` + PageSize int `json:"page_size"` +} + +type OpsIngressRejectHealth struct { + Cardinality int64 `json:"cardinality"` + Capacity int `json:"capacity"` + PendingBatches int `json:"pending_batches"` + PendingRows int `json:"pending_rows"` + Overflowed uint64 `json:"overflowed_count"` + Dropped uint64 `json:"dropped_count"` + Flushed uint64 `json:"flushed_request_count"` + FlushFailures uint64 `json:"flush_failure_count"` + Accepting bool `json:"accepting"` + LastError string `json:"last_error,omitempty"` +} + +type OpsIngressRejectRepository interface { + BatchUpsertIngressRejects(ctx context.Context, items []*OpsIngressRejectAggregate) error + ListIngressRejects(ctx context.Context, filter *OpsIngressRejectFilter) (*OpsIngressRejectList, error) +} + +type ingressRejectKey struct { + reason string + routeFamily string + protocol string + clientIP string + userID int64 + apiKeyID int64 +} + +type ingressRejectShard struct { + mu sync.Mutex + items map[ingressRejectKey]*OpsIngressRejectAggregate +} + +type OpsIngressRejectAggregator struct { + repo OpsIngressRejectRepository + + shards [ingressRejectShardCount]ingressRejectShard + recordMu sync.RWMutex + snapshotMu sync.Mutex + bucket atomic.Int64 + cardinality atomic.Int64 + overflowMu sync.Mutex + overflow *OpsIngressRejectAggregate + + pendingMu sync.Mutex + pending [][]*OpsIngressRejectAggregate + + ctx context.Context + cancel context.CancelFunc + wg sync.WaitGroup + flushCh chan struct{} + started atomic.Bool + accepting atomic.Bool + stopOnce sync.Once + + overflowed atomic.Uint64 + dropped atomic.Uint64 + flushed atomic.Uint64 + flushFailures atomic.Uint64 + lastError atomic.Value +} + +func NewOpsIngressRejectAggregator(repo OpsIngressRejectRepository) *OpsIngressRejectAggregator { + ctx, cancel := context.WithCancel(context.Background()) + a := &OpsIngressRejectAggregator{ + repo: repo, ctx: ctx, cancel: cancel, flushCh: make(chan struct{}, 1), + } + for i := range a.shards { + a.shards[i].items = make(map[ingressRejectKey]*OpsIngressRejectAggregate) + } + a.lastError.Store("") + return a +} + +func (a *OpsIngressRejectAggregator) Start() { + if a == nil || a.repo == nil || !a.started.CompareAndSwap(false, true) { + return + } + a.accepting.Store(true) + a.wg.Add(1) + go func() { + defer a.wg.Done() + ticker := time.NewTicker(ingressRejectFlushInterval) + defer ticker.Stop() + for { + select { + case <-a.ctx.Done(): + return + case <-ticker.C: + a.snapshotAndEnqueue(false) + a.flushPending() + case <-a.flushCh: + a.flushPending() + } + } + }() +} + +func (a *OpsIngressRejectAggregator) Stop() { + if a == nil { + return + } + a.stopOnce.Do(func() { + a.accepting.Store(false) + a.recordMu.Lock() + a.cancel() + a.recordMu.Unlock() + a.wg.Wait() + a.snapshotAndEnqueue(false) + a.flushPending() + }) +} + +func (a *OpsIngressRejectAggregator) RecordIngressReject(reason, routeFamily, protocol, clientIP string, userID, apiKeyID int64) { + if a == nil || a.repo == nil || !a.accepting.Load() { + return + } + now := time.Now().UTC() + bucket := now.Truncate(ingressRejectBucketSize) + key := ingressRejectKey{ + reason: boundedDimension(reason, "unknown"), routeFamily: boundedDimension(routeFamily, "other"), + protocol: boundedDimension(protocol, "other"), clientIP: boundedDimension(clientIP, "0.0.0.0"), + userID: userID, apiKeyID: apiKeyID, + } + for { + a.ensureBucket(bucket) + a.recordMu.RLock() + if !a.accepting.Load() { + a.recordMu.RUnlock() + return + } + bucketUnix := a.bucket.Load() + if bucketUnix != bucket.Unix() { + a.recordMu.RUnlock() + continue + } + shard := &a.shards[ingressRejectHash(key)%ingressRejectShardCount] + shard.mu.Lock() + if item := shard.items[key]; item != nil { + item.RequestCount++ + item.LastSeen = now + shard.mu.Unlock() + a.recordMu.RUnlock() + return + } + if a.reserveDimension() { + shard.items[key] = aggregateFromKey(key, bucket, now) + shard.mu.Unlock() + a.recordMu.RUnlock() + return + } + shard.mu.Unlock() + a.recordOverflow(bucket, now) + a.recordMu.RUnlock() + return + } +} + +func (a *OpsIngressRejectAggregator) ensureBucket(bucket time.Time) { + if a.bucket.Load() == bucket.Unix() { + return + } + a.snapshotMu.Lock() + defer a.snapshotMu.Unlock() + a.recordMu.Lock() + defer a.recordMu.Unlock() + if a.bucket.Load() == bucket.Unix() { + return + } + items := a.snapshotLocked(true) + a.enqueue(items) + a.bucket.Store(bucket.Unix()) + a.cardinality.Store(0) + select { + case a.flushCh <- struct{}{}: + default: + } +} + +func (a *OpsIngressRejectAggregator) reserveDimension() bool { + for { + current := a.cardinality.Load() + if current >= ingressRejectMaxEntries-1 { + return false + } + if a.cardinality.CompareAndSwap(current, current+1) { + return true + } + } +} + +func (a *OpsIngressRejectAggregator) recordOverflow(bucket, now time.Time) { + a.overflowed.Add(1) + a.overflowMu.Lock() + defer a.overflowMu.Unlock() + if a.overflow == nil || !a.overflow.BucketStart.Equal(bucket) { + a.overflow = &OpsIngressRejectAggregate{ + BucketStart: bucket, RejectReason: "other", RouteFamily: "other", Protocol: "other", + ClientIP: "0.0.0.0", RequestCount: 1, FirstSeen: now, LastSeen: now, + } + return + } + a.overflow.RequestCount++ + a.overflow.LastSeen = now +} + +func (a *OpsIngressRejectAggregator) snapshotAndEnqueue(reset bool) { + if a == nil || a.repo == nil { + return + } + a.snapshotMu.Lock() + items := a.snapshotLocked(reset) + a.snapshotMu.Unlock() + a.enqueue(items) +} + +// snapshotLocked captures counter deltas. When reset is false the dimension keys +// remain resident for the whole minute, so periodic flushes cannot reset the budget. +func (a *OpsIngressRejectAggregator) snapshotLocked(reset bool) []*OpsIngressRejectAggregate { + items := make([]*OpsIngressRejectAggregate, 0) + for i := range a.shards { + shard := &a.shards[i] + shard.mu.Lock() + for key, item := range shard.items { + if item.RequestCount > 0 { + copyItem := *item + items = append(items, ©Item) + item.RequestCount = 0 + } + if reset { + delete(shard.items, key) + } + } + shard.mu.Unlock() + } + a.overflowMu.Lock() + if a.overflow != nil && a.overflow.RequestCount > 0 { + copyItem := *a.overflow + items = append(items, ©Item) + a.overflow.RequestCount = 0 + } + if reset { + a.overflow = nil + } + a.overflowMu.Unlock() + return items +} + +func (a *OpsIngressRejectAggregator) enqueue(items []*OpsIngressRejectAggregate) { + if len(items) == 0 { + return + } + a.pendingMu.Lock() + defer a.pendingMu.Unlock() + if len(a.pending) >= ingressRejectMaxPendingBatches { + a.dropped.Add(batchRequestCount(items)) + return + } + a.pending = append(a.pending, items) +} + +func (a *OpsIngressRejectAggregator) flushPending() { + for { + a.pendingMu.Lock() + if len(a.pending) == 0 { + a.pendingMu.Unlock() + return + } + batch := a.pending[0] + a.pendingMu.Unlock() + + ctx, cancel := context.WithTimeout(context.Background(), ingressRejectFlushTimeout) + err := a.repo.BatchUpsertIngressRejects(ctx, batch) + cancel() + if err != nil { + a.flushFailures.Add(1) + a.lastError.Store(err.Error()) + log.Printf("[IngressRejectAggregator] flush failed: %v", err) + return + } + a.pendingMu.Lock() + if len(a.pending) > 0 { + a.pending = a.pending[1:] + } + a.pendingMu.Unlock() + a.flushed.Add(batchRequestCount(batch)) + a.lastError.Store("") + } +} + +func batchRequestCount(items []*OpsIngressRejectAggregate) uint64 { + var count uint64 + for _, item := range items { + if item != nil && item.RequestCount > 0 { + count += uint64(item.RequestCount) + } + } + return count +} + +func (a *OpsIngressRejectAggregator) Health() OpsIngressRejectHealth { + h := OpsIngressRejectHealth{Capacity: ingressRejectMaxEntries} + if a == nil { + return h + } + h.Cardinality = a.cardinality.Load() + a.overflowMu.Lock() + if a.overflow != nil { + h.Cardinality++ + } + a.overflowMu.Unlock() + a.pendingMu.Lock() + h.PendingBatches = len(a.pending) + for _, batch := range a.pending { + h.PendingRows += len(batch) + } + a.pendingMu.Unlock() + h.Overflowed = a.overflowed.Load() + h.Dropped = a.dropped.Load() + h.Flushed = a.flushed.Load() + h.FlushFailures = a.flushFailures.Load() + h.Accepting = a.accepting.Load() + if v := a.lastError.Load(); v != nil { + h.LastError, _ = v.(string) + } + return h +} + +func boundedDimension(value, fallback string) string { + value = strings.TrimSpace(value) + if value == "" { + return fallback + } + if len(value) > 64 { + return value[:64] + } + return value +} + +func ingressRejectHash(k ingressRejectKey) int { + h := uint64(1469598103934665603) + for _, value := range []string{k.reason, k.routeFamily, k.protocol, k.clientIP} { + for i := 0; i < len(value); i++ { + h ^= uint64(value[i]) + h *= 1099511628211 + } + } + h ^= uint64(k.userID) + h *= 1099511628211 + h ^= uint64(k.apiKeyID) + return int(h & 0x7fffffff) +} + +func aggregateFromKey(k ingressRejectKey, bucket, now time.Time) *OpsIngressRejectAggregate { + item := &OpsIngressRejectAggregate{ + BucketStart: bucket, RejectReason: k.reason, RouteFamily: k.routeFamily, Protocol: k.protocol, + ClientIP: k.clientIP, RequestCount: 1, FirstSeen: now, LastSeen: now, + } + if k.userID > 0 { + value := k.userID + item.UserID = &value + } + if k.apiKeyID > 0 { + value := k.apiKeyID + item.APIKeyID = &value + } + return item +} + +func (s *OpsService) SetIngressRejectAggregator(a *OpsIngressRejectAggregator) { + if s != nil { + s.ingressRejectAggregator = a + } +} + +func (s *OpsService) RecordIngressReject(reason, routeFamily, protocol, clientIP string, userID, apiKeyID int64) { + if s != nil && s.ingressRejectAggregator != nil { + s.ingressRejectAggregator.RecordIngressReject(reason, routeFamily, protocol, clientIP, userID, apiKeyID) + } +} + +func (s *OpsService) GetIngressRejectHealth() OpsIngressRejectHealth { + if s == nil || s.ingressRejectAggregator == nil { + return OpsIngressRejectHealth{Capacity: ingressRejectMaxEntries} + } + return s.ingressRejectAggregator.Health() +} + +func (s *OpsService) ListIngressRejects(ctx context.Context, filter *OpsIngressRejectFilter) (*OpsIngressRejectList, error) { + if err := s.RequireMonitoringEnabled(ctx); err != nil { + return nil, err + } + repo, ok := s.opsRepo.(OpsIngressRejectRepository) + if !ok { + return &OpsIngressRejectList{Items: []*OpsIngressRejectAggregate{}, Page: 1, PageSize: 50}, nil + } + return repo.ListIngressRejects(ctx, filter) +} diff --git a/backend/internal/service/ops_ingress_reject_test.go b/backend/internal/service/ops_ingress_reject_test.go new file mode 100644 index 000000000..6a67aafbf --- /dev/null +++ b/backend/internal/service/ops_ingress_reject_test.go @@ -0,0 +1,115 @@ +package service + +import ( + "context" + "errors" + "fmt" + "sync" + "testing" + + "github.com/stretchr/testify/require" +) + +type ingressRejectRepoStub struct { + mu sync.Mutex + failCount int + calls int + requests int64 +} + +func (r *ingressRejectRepoStub) BatchUpsertIngressRejects(_ context.Context, items []*OpsIngressRejectAggregate) error { + r.mu.Lock() + defer r.mu.Unlock() + r.calls++ + if r.failCount > 0 { + r.failCount-- + return errors.New("temporary database failure") + } + for _, item := range items { + if item != nil { + r.requests += item.RequestCount + } + } + return nil +} + +func (r *ingressRejectRepoStub) ListIngressRejects(context.Context, *OpsIngressRejectFilter) (*OpsIngressRejectList, error) { + return &OpsIngressRejectList{}, nil +} + +func TestOpsIngressRejectAggregatorUsesGlobalBucketCapacity(t *testing.T) { + repo := &ingressRejectRepoStub{} + a := NewOpsIngressRejectAggregator(repo) + a.Start() + + // Concentrating all dimensions in one shard must not waste capacity in the others. + inserted := 0 + for i := 0; inserted < 600; i++ { + ip := fmt.Sprintf("target-%d", i) + key := ingressRejectKey{reason: "invalid_api_key", routeFamily: "messages", protocol: "anthropic", clientIP: ip} + if ingressRejectHash(key)%ingressRejectShardCount == 0 { + a.RecordIngressReject(key.reason, key.routeFamily, key.protocol, ip, 0, 0) + inserted++ + } + } + require.Equal(t, int64(600), a.Health().Cardinality) + + for i := 0; i < ingressRejectMaxEntries; i++ { + a.RecordIngressReject("invalid_api_key", "messages", "anthropic", fmt.Sprintf("rotating-%d", i), 0, 0) + } + health := a.Health() + require.Equal(t, int64(ingressRejectMaxEntries), health.Cardinality) + require.Greater(t, health.Overflowed, uint64(0)) + + a.snapshotAndEnqueue(false) + require.Equal(t, int64(ingressRejectMaxEntries), a.Health().Cardinality, "periodic flush must retain the minute budget") + a.Stop() +} + +func TestOpsIngressRejectAggregatorConcurrentCountAndStopFlush(t *testing.T) { + repo := &ingressRejectRepoStub{} + a := NewOpsIngressRejectAggregator(repo) + a.Start() + + const goroutines = 32 + const perGoroutine = 200 + var wg sync.WaitGroup + for i := 0; i < goroutines; i++ { + wg.Add(1) + go func() { + defer wg.Done() + for j := 0; j < perGoroutine; j++ { + a.RecordIngressReject("invalid_api_key", "responses", "openai", "192.0.2.10", 0, 0) + } + }() + } + wg.Wait() + a.Stop() + + repo.mu.Lock() + require.Equal(t, int64(goroutines*perGoroutine), repo.requests) + repo.mu.Unlock() + require.False(t, a.Health().Accepting) + a.RecordIngressReject("invalid_api_key", "responses", "openai", "192.0.2.10", 0, 0) +} + +func TestOpsIngressRejectAggregatorRetriesBoundedPendingBatch(t *testing.T) { + repo := &ingressRejectRepoStub{failCount: 1} + a := NewOpsIngressRejectAggregator(repo) + a.accepting.Store(true) + a.RecordIngressReject("group_deleted", "messages", "anthropic", "192.0.2.20", 1, 2) + a.snapshotAndEnqueue(false) + a.flushPending() + health := a.Health() + require.Equal(t, 1, health.PendingBatches) + require.Equal(t, uint64(1), health.FlushFailures) + require.Equal(t, uint64(0), health.Dropped) + + a.flushPending() + require.Equal(t, 0, a.Health().PendingBatches) + a.Stop() + repo.mu.Lock() + require.Equal(t, int64(1), repo.requests) + require.GreaterOrEqual(t, repo.calls, 2) + repo.mu.Unlock() +} diff --git a/backend/internal/service/ops_log_runtime_test.go b/backend/internal/service/ops_log_runtime_test.go index 658b48128..0a8cfa906 100644 --- a/backend/internal/service/ops_log_runtime_test.go +++ b/backend/internal/service/ops_log_runtime_test.go @@ -11,12 +11,14 @@ import ( ) type runtimeSettingRepoStub struct { - values map[string]string - deleted map[string]bool - setCalls int - getValueFn func(key string) (string, error) - setFn func(key, value string) error - deleteFn func(key string) error + values map[string]string + deleted map[string]bool + setCalls int + getValueCalls int + getMultipleCalls int + getValueFn func(key string) (string, error) + setFn func(key, value string) error + deleteFn func(key string) error } func newRuntimeSettingRepoStub() *runtimeSettingRepoStub { @@ -35,6 +37,7 @@ func (s *runtimeSettingRepoStub) Get(ctx context.Context, key string) (*Setting, } func (s *runtimeSettingRepoStub) GetValue(_ context.Context, key string) (string, error) { + s.getValueCalls++ if s.getValueFn != nil { return s.getValueFn(key) } @@ -57,6 +60,7 @@ func (s *runtimeSettingRepoStub) Set(_ context.Context, key, value string) error } func (s *runtimeSettingRepoStub) GetMultiple(_ context.Context, keys []string) (map[string]string, error) { + s.getMultipleCalls++ out := make(map[string]string, len(keys)) for _, key := range keys { if value, ok := s.values[key]; ok { diff --git a/backend/internal/service/ops_models.go b/backend/internal/service/ops_models.go index 411ca7595..a6830cadc 100644 --- a/backend/internal/service/ops_models.go +++ b/backend/internal/service/ops_models.go @@ -74,11 +74,6 @@ type OpsErrorLog struct { // 关联 api_key 名称(LEFT JOIN api_keys 取得;软删只覆盖 key 列,name 保留,故已删 key 仍有原名)。 APIKeyName string `json:"api_key_name,omitempty"` APIKeyDeleted bool `json:"api_key_deleted,omitempty"` - - // 已删除 KEY 所有者(INVALID_API_KEY 且该 key 曾存在时的归因快照)。 - // 认证失败行 user_id 为空,列表用户列以此回退显示所有者。 - DeletedKeyOwnerUserID *int64 `json:"deleted_key_owner_user_id,omitempty"` - DeletedKeyOwnerEmail string `json:"deleted_key_owner_email,omitempty"` } type OpsErrorLogDetail struct { @@ -102,12 +97,7 @@ type OpsErrorLogDetail struct { // vNext metric semantics IsBusinessLimited bool `json:"is_business_limited"` - // Deleted key owner info (populated when INVALID_API_KEY and key was previously deleted). - // OwnerUserID/OwnerEmail 已上移到 OpsErrorLog(列表用户列回退需要)。 - AttemptedKeyPrefix string `json:"attempted_key_prefix,omitempty"` - DeletedKeyName string `json:"deleted_key_name,omitempty"` - - // Bound (non-deleted) key prefix, snapshotted at error time; mutually exclusive with AttemptedKeyPrefix. + // Bound (non-deleted) key prefix, snapshotted at error time. APIKeyPrefix string `json:"api_key_prefix,omitempty"` } @@ -137,11 +127,6 @@ type OpsErrorLogFilter struct { UserID *int64 APIKeyID *int64 - // MatchDeletedKeyOwner: 用户侧专用。UserID 设置且为 true 时,归属从 user_id=UserID - // 放宽为 (user_id=UserID OR deleted_key_owner_user_id=UserID),使原所有者能看到 - // 自己「已删除 key 认证失败」的记录。admin 路径不设此开关 → 行为不变。 - MatchDeletedKeyOwner bool - // Model matches against requested_model first, then model. Model string // ModelFuzzy 为 true 时 Model 走 ILIKE 模糊匹配(仅用户端启用);false(默认)保持精确 =,管理端语义不变。 diff --git a/backend/internal/service/ops_port.go b/backend/internal/service/ops_port.go index 2b73d4a69..02186b6a2 100644 --- a/backend/internal/service/ops_port.go +++ b/backend/internal/service/ops_port.go @@ -10,8 +10,6 @@ type OpsRepository interface { BatchInsertErrorLogs(ctx context.Context, inputs []*OpsInsertErrorLogInput) (int64, error) ListErrorLogs(ctx context.Context, filter *OpsErrorLogFilter) (*OpsErrorLogList, error) GetErrorLogByID(ctx context.Context, id int64) (*OpsErrorLogDetail, error) - // LookupDeletedKeyAudit 按明文 key 反查最近一条已删除 key 审计;未命中返回 (nil, nil)。 - LookupDeletedKeyAudit(ctx context.Context, key string) (*DeletedKeyAuditResult, error) ListRequestDetails(ctx context.Context, filter *OpsRequestDetailFilter) ([]*OpsRequestDetail, int64, error) BatchInsertSystemLogs(ctx context.Context, inputs []*OpsInsertSystemLogInput) (int64, error) ListSystemLogs(ctx context.Context, filter *OpsSystemLogFilter) (*OpsSystemLogList, error) @@ -63,12 +61,6 @@ type OpsRepository interface { GetLatestDailyBucketDate(ctx context.Context) (time.Time, bool, error) } -// DeletedKeyAuditResult 是按明文 key 反查 deleted_api_key_audits 的结果。 -type DeletedKeyAuditResult struct { - UserID int64 - KeyName string -} - type OpsInsertErrorLogInput struct { RequestID string ClientRequestID string @@ -127,12 +119,7 @@ type OpsInsertErrorLogInput struct { CreatedAt time.Time - // 已删除 key 归因(仅 INVALID_API_KEY 认证失败时可能非空) - AttemptedKeyPrefix string // 提交 key 的脱敏前缀(前 8 位) - DeletedKeyOwnerUserID *int64 // 反查命中的原所有者 user_id - DeletedKeyName string // 反查命中的 key 名称 - - // 有效(未删除)key 报错时快照的 key 脱敏前缀(前 8 位);与 AttemptedKeyPrefix 互斥。 + // 有效(未删除)key 报错时快照的 key 脱敏前缀(前 8 位)。 // 落库快照而非读时 JOIN:key 之后被删(key 列被 tombstone 覆盖)仍保留当时前缀。 APIKeyPrefix string } diff --git a/backend/internal/service/ops_queue_sanitize_test.go b/backend/internal/service/ops_queue_sanitize_test.go new file mode 100644 index 000000000..f2ae621c0 --- /dev/null +++ b/backend/internal/service/ops_queue_sanitize_test.go @@ -0,0 +1,49 @@ +package service + +import ( + "strings" + "testing" +) + +func TestSanitizeOpsUpstreamErrorsForQueueBoundsAndRedacts(t *testing.T) { + entry := &OpsInsertErrorLogInput{} + for i := 0; i < 20; i++ { + entry.UpstreamErrors = append(entry.UpstreamErrors, &OpsUpstreamErrorEvent{ + Platform: strings.Repeat("p", 100), + AccountName: strings.Repeat("a", 300), + UpstreamStatusCode: 500, + UpstreamURL: strings.Repeat("u", 3000), + UpstreamResponseBody: `{"authorization":"Bearer secret","message":"` + strings.Repeat("x", 10_000) + `"}`, + Message: strings.Repeat("m", 3000), + Detail: `{"api_key":"secret","detail":"` + strings.Repeat("y", 10_000) + `"}`, + }) + } + + if err := SanitizeOpsUpstreamErrorsForQueue(entry); err != nil { + t.Fatal(err) + } + if entry.UpstreamErrors != nil { + t.Fatal("raw upstream event slice must be released before queueing") + } + if entry.UpstreamErrorsJSON == nil { + t.Fatal("sanitized upstream event JSON is missing") + } + events, err := ParseOpsUpstreamErrors(*entry.UpstreamErrorsJSON) + if err != nil { + t.Fatal(err) + } + if len(events) != 16 { + t.Fatalf("event count = %d, want 16", len(events)) + } + for _, event := range events { + if len(event.Platform) > 32 || len(event.AccountName) > 128 || len(event.UpstreamURL) > 2048 || len(event.Message) > 2048 { + t.Fatalf("event fields were not bounded: %+v", event) + } + if len(event.UpstreamResponseBody) > OpsErrorLogQueueBodyMaxBytes || len(event.Detail) > OpsErrorLogQueueBodyMaxBytes { + t.Fatal("event body/detail exceeded queue limit") + } + if strings.Contains(event.UpstreamResponseBody, "Bearer secret") || strings.Contains(event.Detail, `"secret"`) { + t.Fatal("credential material was not redacted") + } + } +} diff --git a/backend/internal/service/ops_repo_mock_test.go b/backend/internal/service/ops_repo_mock_test.go index 5e33bffb1..4138ea77a 100644 --- a/backend/internal/service/ops_repo_mock_test.go +++ b/backend/internal/service/ops_repo_mock_test.go @@ -13,7 +13,6 @@ type opsRepoMock struct { ListSystemLogsFn func(ctx context.Context, filter *OpsSystemLogFilter) (*OpsSystemLogList, error) DeleteSystemLogsFn func(ctx context.Context, filter *OpsSystemLogCleanupFilter) (int64, error) InsertSystemLogCleanupAuditFn func(ctx context.Context, input *OpsSystemLogCleanupAudit) error - LookupDeletedKeyAuditFn func(ctx context.Context, key string) (*DeletedKeyAuditResult, error) } func (m *opsRepoMock) InsertErrorLog(ctx context.Context, input *OpsInsertErrorLogInput) (int64, error) { @@ -190,11 +189,4 @@ func (m *opsRepoMock) GetLatestDailyBucketDate(ctx context.Context) (time.Time, return time.Time{}, false, nil } -func (m *opsRepoMock) LookupDeletedKeyAudit(ctx context.Context, key string) (*DeletedKeyAuditResult, error) { - if m.LookupDeletedKeyAuditFn != nil { - return m.LookupDeletedKeyAuditFn(ctx, key) - } - return nil, nil -} - var _ OpsRepository = (*opsRepoMock)(nil) diff --git a/backend/internal/service/ops_runtime_snapshot_test.go b/backend/internal/service/ops_runtime_snapshot_test.go new file mode 100644 index 000000000..9e0631ecb --- /dev/null +++ b/backend/internal/service/ops_runtime_snapshot_test.go @@ -0,0 +1,186 @@ +package service + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "testing" + "time" +) + +var ( + benchmarkOpsMonitoringEnabled bool + benchmarkOpsAdvancedSettings OpsAdvancedSettings +) + +type opsRuntimeRefreshRepo struct { + SettingRepository + mu sync.RWMutex + values map[string]string + fail atomic.Bool + calls atomic.Int64 +} + +func (r *opsRuntimeRefreshRepo) GetMultiple(ctx context.Context, keys []string) (map[string]string, error) { + r.calls.Add(1) + select { + case <-ctx.Done(): + return nil, ctx.Err() + default: + } + if r.fail.Load() { + return nil, errors.New("settings unavailable") + } + r.mu.RLock() + defer r.mu.RUnlock() + out := make(map[string]string, len(keys)) + for _, key := range keys { + if value, ok := r.values[key]; ok { + out[key] = value + } + } + return out, nil +} + +func (r *opsRuntimeRefreshRepo) set(key, value string) { + r.mu.Lock() + r.values[key] = value + r.mu.Unlock() +} + +func waitForOpsRefresh(t *testing.T, timeout time.Duration, condition func() bool) { + t.Helper() + deadline := time.Now().Add(timeout) + for time.Now().Before(deadline) { + if condition() { + return + } + time.Sleep(time.Millisecond) + } + t.Fatal("condition not satisfied before timeout") +} + +func TestOpsRuntimeSettingsSnapshotLoadsOnceAndServesHotPath(t *testing.T) { + repo := newRuntimeSettingRepoStub() + repo.values[SettingKeyOpsMonitoringEnabled] = "false" + repo.values[SettingKeyOpsAdvancedSettings] = `{"ignore_context_canceled":false,"auto_refresh_interval_seconds":45}` + + svc := &OpsService{settingRepo: repo} + svc.initRuntimeSettings(context.Background()) + if repo.getMultipleCalls != 1 { + t.Fatalf("startup GetMultiple calls = %d, want 1", repo.getMultipleCalls) + } + + for range 1000 { + if svc.IsMonitoringEnabled(context.Background()) { + t.Fatal("monitoring enabled, want false") + } + cfg, err := svc.GetOpsAdvancedSettings(context.Background()) + if err != nil { + t.Fatalf("GetOpsAdvancedSettings() error = %v", err) + } + if cfg.AutoRefreshIntervalSec != 45 { + t.Fatalf("AutoRefreshIntervalSec = %d, want 45", cfg.AutoRefreshIntervalSec) + } + } + if repo.getValueCalls != 0 || repo.getMultipleCalls != 1 { + t.Fatalf("hot path touched repository: get=%d get_multiple=%d", repo.getValueCalls, repo.getMultipleCalls) + } +} + +func TestOpsRuntimeSettingsAdministrativeUpdatesAreImmediatelyVisible(t *testing.T) { + svc := &OpsService{} + svc.initRuntimeSettings(context.Background()) + + svc.SetMonitoringEnabled(false) + if svc.IsMonitoringEnabled(context.Background()) { + t.Fatal("monitoring update was not visible") + } + + cfg := defaultOpsAdvancedSettings() + cfg.IgnoreNoAvailableAccounts = true + svc.storeAdvancedSettingsSnapshot(cfg) + got, err := svc.GetOpsAdvancedSettings(context.Background()) + if err != nil { + t.Fatalf("GetOpsAdvancedSettings() error = %v", err) + } + if !got.IgnoreNoAvailableAccounts { + t.Fatal("advanced settings update was not visible") + } + if svc.IsMonitoringEnabled(context.Background()) { + t.Fatal("advanced update overwrote monitoring setting") + } +} + +func TestOpsRuntimeSettingsBackgroundRefreshConverges(t *testing.T) { + repo := &opsRuntimeRefreshRepo{values: map[string]string{SettingKeyOpsMonitoringEnabled: "false"}} + svc := &OpsService{settingRepo: repo} + svc.initRuntimeSettings(context.Background()) + if svc.IsMonitoringEnabled(context.Background()) { + t.Fatal("initial monitoring state = true, want false") + } + + repo.set(SettingKeyOpsMonitoringEnabled, "true") + svc.startRuntimeSettingsRefresh(context.Background(), 5*time.Millisecond, 0, 50*time.Millisecond) + t.Cleanup(svc.StopRuntimeSettingsRefresh) + waitForOpsRefresh(t, time.Second, func() bool { + return svc.IsMonitoringEnabled(context.Background()) && svc.RuntimeSettingsRefreshHealth().SuccessTotal > 0 + }) +} + +func TestOpsRuntimeSettingsRefreshFailuresKeepLastKnownGoodSnapshot(t *testing.T) { + repo := &opsRuntimeRefreshRepo{values: map[string]string{ + SettingKeyOpsMonitoringEnabled: "false", + SettingKeyOpsAdvancedSettings: `{"ignore_no_available_accounts":true}`, + }} + svc := &OpsService{settingRepo: repo} + svc.initRuntimeSettings(context.Background()) + repo.fail.Store(true) + svc.startRuntimeSettingsRefresh(context.Background(), 5*time.Millisecond, 0, 50*time.Millisecond) + t.Cleanup(svc.StopRuntimeSettingsRefresh) + waitForOpsRefresh(t, time.Second, func() bool { + return svc.RuntimeSettingsRefreshHealth().FailureTotal >= 3 + }) + + if svc.IsMonitoringEnabled(context.Background()) { + t.Fatal("failed refresh overwrote last known monitoring state") + } + if !svc.OpsAdvancedSettingsSnapshot().IgnoreNoAvailableAccounts { + t.Fatal("failed refresh overwrote last known advanced settings") + } +} + +func TestOpsRuntimeSettingsRefreshStopEndsLifecycle(t *testing.T) { + repo := &opsRuntimeRefreshRepo{values: map[string]string{}} + svc := &OpsService{settingRepo: repo} + svc.initRuntimeSettings(context.Background()) + svc.startRuntimeSettingsRefresh(context.Background(), 5*time.Millisecond, 0, 50*time.Millisecond) + waitForOpsRefresh(t, time.Second, func() bool { + return svc.RuntimeSettingsRefreshHealth().SuccessTotal > 0 + }) + + svc.StopRuntimeSettingsRefresh() + callsAfterStop := repo.calls.Load() + time.Sleep(20 * time.Millisecond) + if got := repo.calls.Load(); got != callsAfterStop { + t.Fatalf("refresh continued after Stop: before=%d after=%d", callsAfterStop, got) + } + if svc.RuntimeSettingsRefreshHealth().Running { + t.Fatal("refresh health still reports running after Stop") + } + // Idempotence is part of the cleanup contract. + svc.StopRuntimeSettingsRefresh() +} + +func BenchmarkOpsRuntimeSettingsSnapshotRead(b *testing.B) { + svc := &OpsService{} + svc.initRuntimeSettings(context.Background()) + ctx := context.Background() + b.ReportAllocs() + b.ResetTimer() + for range b.N { + benchmarkOpsMonitoringEnabled = svc.IsMonitoringEnabled(ctx) + benchmarkOpsAdvancedSettings = svc.OpsAdvancedSettingsSnapshot() + } +} diff --git a/backend/internal/service/ops_service.go b/backend/internal/service/ops_service.go index 3b2725266..b1d9b9f57 100644 --- a/backend/internal/service/ops_service.go +++ b/backend/internal/service/ops_service.go @@ -6,7 +6,10 @@ import ( "encoding/json" "errors" "log" + "math/rand/v2" "strings" + "sync" + "sync/atomic" "time" "github.com/Wei-Shaw/sub2api/internal/config" @@ -17,8 +20,27 @@ var ErrOpsDisabled = infraerrors.NotFound("OPS_DISABLED", "Ops monitoring is dis const ( opsMaxStoredErrorBodyBytes = 20 * 1024 + // OpsErrorLogQueueBodyMaxBytes bounds attacker-controlled response data while + // it waits in the asynchronous error-log queue. + OpsErrorLogQueueBodyMaxBytes = 8 * 1024 + + opsRuntimeSettingsRefreshInterval = 30 * time.Second + opsRuntimeSettingsRefreshJitter = 20 + opsRuntimeSettingsRefreshTimeout = 3 * time.Second + opsRuntimeSettingsFailureLogEvery = time.Minute ) +type opsRuntimeSettingsSnapshot struct { + monitoringEnabled bool + advanced OpsAdvancedSettings +} + +type OpsRuntimeSettingsRefreshHealth struct { + Running bool `json:"running"` + SuccessTotal uint64 `json:"success_total"` + FailureTotal uint64 `json:"failure_total"` +} + // OpsService provides ingestion and query APIs for the Ops monitoring module. type OpsService struct { opsRepo OpsRepository @@ -31,12 +53,15 @@ type OpsService struct { // getAccountAvailability is a unit-test hook for overriding account availability lookup. getAccountAvailability func(ctx context.Context, platformFilter string, groupIDFilter *int64) (*OpsAccountAvailability, error) - concurrencyService *ConcurrencyService - gatewayService *GatewayService - openAIGatewayService *OpenAIGatewayService - geminiCompatService *GeminiMessagesCompatService - antigravityGatewayService *AntigravityGatewayService - systemLogSink *OpsSystemLogSink + concurrencyService *ConcurrencyService + gatewayService *GatewayService + openAIGatewayService *OpenAIGatewayService + geminiCompatService *GeminiMessagesCompatService + antigravityGatewayService *AntigravityGatewayService + systemLogSink *OpsSystemLogSink + ingressRejectAggregator *OpsIngressRejectAggregator + authCacheInvalidationWorker *AuthCacheInvalidationWorker + apiKeyService *APIKeyService // cleanupReloader 由 wire 在 OpsCleanupService 构造完成后通过 SetCleanupReloader 注入。 // 解耦避免 OpsService -> OpsCleanupService 的硬依赖(cleanup 也读 settings,会循环)。 @@ -46,6 +71,19 @@ type OpsService struct { // UpdateOpsAdvancedSettings 写入新配置后调用,把最新的 quota auto-pause 全局默认阈值 // 立即同步到调度热路径读取的内存缓存,避免下次请求才能感知新值。 quotaAutoPauseSink func(OpsOpenAIAccountQuotaAutoPauseSettings) + + // Published snapshots are immutable. Gateway reads are lock-free; the mutex + // only serializes startup and administrative updates. + runtimeSettings atomic.Pointer[opsRuntimeSettingsSnapshot] + runtimeSettingsMu sync.Mutex + + runtimeRefreshMu sync.Mutex + runtimeRefreshCancel context.CancelFunc + runtimeRefreshDone chan struct{} + runtimeRefreshRunning atomic.Bool + runtimeRefreshSuccess atomic.Uint64 + runtimeRefreshFailure atomic.Uint64 + runtimeRefreshLastFailureLog atomic.Int64 } // CleanupReloader 由 OpsCleanupService 实现。 @@ -100,6 +138,7 @@ func NewOpsService( antigravityGatewayService: antigravityGatewayService, systemLogSink: systemLogSink, } + svc.initRuntimeSettings(context.Background()) svc.applyRuntimeLogConfigOnStartup(context.Background()) return svc } @@ -112,22 +151,20 @@ func (s *OpsService) RequireMonitoringEnabled(ctx context.Context) error { } func (s *OpsService) IsMonitoringEnabled(ctx context.Context) bool { + _ = ctx // Hard switch: disable ops entirely. if s.cfg != nil && !s.cfg.Ops.Enabled { return false } - if s.settingRepo == nil { - return true - } - value, err := s.settingRepo.GetValue(ctx, SettingKeyOpsMonitoringEnabled) - if err != nil { - // Default enabled when key is missing, and fail-open on transient errors - // (ops should never block gateway traffic). - if errors.Is(err, ErrSettingNotFound) { - return true - } - return true + if snapshot := s.runtimeSettings.Load(); snapshot != nil { + return snapshot.monitoringEnabled } + // Directly assembled test services and failed cold loads remain fail-open, + // without turning a request into a settings-table lookup. + return true +} + +func parseOpsMonitoringEnabled(value string) bool { switch strings.ToLower(strings.TrimSpace(value)) { case "false", "0", "off", "disabled": return false @@ -136,6 +173,227 @@ func (s *OpsService) IsMonitoringEnabled(ctx context.Context) bool { } } +func (s *OpsService) initRuntimeSettings(ctx context.Context) { + if s == nil { + return + } + defaults := defaultOpsAdvancedSettings() + s.runtimeSettings.Store(&opsRuntimeSettingsSnapshot{monitoringEnabled: true, advanced: *defaults}) + _ = s.RefreshRuntimeSettings(ctx) +} + +// RefreshRuntimeSettings is the cold-path database load used at startup and by +// explicit administrative refreshes. Request processing only reads the atomic +// snapshot. +func (s *OpsService) RefreshRuntimeSettings(ctx context.Context) error { + if s == nil || s.settingRepo == nil { + return nil + } + if ctx == nil { + ctx = context.Background() + } + s.runtimeSettingsMu.Lock() + defer s.runtimeSettingsMu.Unlock() + + values, err := s.settingRepo.GetMultiple(ctx, []string{ + SettingKeyOpsMonitoringEnabled, + SettingKeyOpsAdvancedSettings, + }) + if err != nil { + return err + } + + monitoringEnabled := true + if raw, ok := values[SettingKeyOpsMonitoringEnabled]; ok { + monitoringEnabled = parseOpsMonitoringEnabled(raw) + } + advanced := defaultOpsAdvancedSettings() + if raw, ok := values[SettingKeyOpsAdvancedSettings]; ok { + if err := json.Unmarshal([]byte(raw), advanced); err != nil { + advanced = defaultOpsAdvancedSettings() + } + } + normalizeOpsAdvancedSettings(advanced) + + s.runtimeSettings.Store(&opsRuntimeSettingsSnapshot{monitoringEnabled: monitoringEnabled, advanced: *advanced}) + return nil +} + +// StartRuntimeSettingsRefresh keeps DB-backed Ops settings converged across +// application instances without putting database I/O on request paths. +func (s *OpsService) StartRuntimeSettingsRefresh(ctx context.Context) { + s.startRuntimeSettingsRefresh(ctx, opsRuntimeSettingsRefreshInterval, opsRuntimeSettingsRefreshJitter, opsRuntimeSettingsRefreshTimeout) +} + +func (s *OpsService) startRuntimeSettingsRefresh(ctx context.Context, interval time.Duration, jitterPercent int, timeout time.Duration) { + if s == nil || s.settingRepo == nil { + return + } + if ctx == nil { + ctx = context.Background() + } + if interval <= 0 { + interval = opsRuntimeSettingsRefreshInterval + } + if timeout <= 0 { + timeout = opsRuntimeSettingsRefreshTimeout + } + if jitterPercent < 0 { + jitterPercent = 0 + } + if jitterPercent > 100 { + jitterPercent = 100 + } + + s.runtimeRefreshMu.Lock() + if s.runtimeRefreshCancel != nil { + s.runtimeRefreshMu.Unlock() + return + } + refreshCtx, cancel := context.WithCancel(ctx) + done := make(chan struct{}) + s.runtimeRefreshCancel = cancel + s.runtimeRefreshDone = done + s.runtimeRefreshRunning.Store(true) + s.runtimeRefreshMu.Unlock() + + go func() { + defer close(done) + defer s.runtimeRefreshRunning.Store(false) + for { + delay := jitterDuration(interval, jitterPercent) + timer := time.NewTimer(delay) + select { + case <-refreshCtx.Done(): + if !timer.Stop() { + select { + case <-timer.C: + default: + } + } + return + case <-timer.C: + } + + attemptCtx, attemptCancel := context.WithTimeout(refreshCtx, timeout) + err := s.RefreshRuntimeSettings(attemptCtx) + attemptCancel() + if err != nil { + s.runtimeRefreshFailure.Add(1) + s.logRuntimeSettingsRefreshFailure(err) + continue + } + s.runtimeRefreshSuccess.Add(1) + } + }() +} + +func jitterDuration(base time.Duration, percent int) time.Duration { + if base <= 0 || percent <= 0 { + return base + } + delta := float64(percent) / 100 + factor := 1 - delta + rand.Float64()*(2*delta) + if factor <= 0 { + return base + } + return time.Duration(float64(base) * factor) +} + +func (s *OpsService) logRuntimeSettingsRefreshFailure(err error) { + if s == nil || err == nil { + return + } + now := time.Now().Unix() + for { + last := s.runtimeRefreshLastFailureLog.Load() + if last != 0 && now-last < int64(opsRuntimeSettingsFailureLogEvery/time.Second) { + return + } + if s.runtimeRefreshLastFailureLog.CompareAndSwap(last, now) { + log.Printf("[Ops] runtime settings refresh failed: %v", err) + return + } + } +} + +// StopRuntimeSettingsRefresh is idempotent and waits for an in-flight refresh +// to observe cancellation before returning. +func (s *OpsService) StopRuntimeSettingsRefresh() { + if s == nil { + return + } + s.runtimeRefreshMu.Lock() + cancel := s.runtimeRefreshCancel + done := s.runtimeRefreshDone + s.runtimeRefreshMu.Unlock() + if cancel != nil { + cancel() + } + if done != nil { + <-done + } + s.runtimeRefreshMu.Lock() + if s.runtimeRefreshDone == done { + s.runtimeRefreshCancel = nil + s.runtimeRefreshDone = nil + } + s.runtimeRefreshMu.Unlock() +} + +func (s *OpsService) RuntimeSettingsRefreshHealth() OpsRuntimeSettingsRefreshHealth { + if s == nil { + return OpsRuntimeSettingsRefreshHealth{} + } + return OpsRuntimeSettingsRefreshHealth{ + Running: s.runtimeRefreshRunning.Load(), + SuccessTotal: s.runtimeRefreshSuccess.Load(), + FailureTotal: s.runtimeRefreshFailure.Load(), + } +} + +// SetMonitoringEnabled publishes an already-persisted admin setting without a +// database round trip. +func (s *OpsService) SetMonitoringEnabled(enabled bool) { + if s == nil { + return + } + s.runtimeSettingsMu.Lock() + current := s.runtimeSettings.Load() + next := &opsRuntimeSettingsSnapshot{monitoringEnabled: enabled, advanced: *defaultOpsAdvancedSettings()} + if current != nil { + next.advanced = current.advanced + } + s.runtimeSettings.Store(next) + s.runtimeSettingsMu.Unlock() +} + +func (s *OpsService) storeAdvancedSettingsSnapshot(cfg *OpsAdvancedSettings) { + if s == nil || cfg == nil { + return + } + s.runtimeSettingsMu.Lock() + current := s.runtimeSettings.Load() + next := &opsRuntimeSettingsSnapshot{monitoringEnabled: true, advanced: *cfg} + if current != nil { + next.monitoringEnabled = current.monitoringEnabled + } + s.runtimeSettings.Store(next) + s.runtimeSettingsMu.Unlock() +} + +// SanitizeOpsErrorBodyForQueue removes credentials and truncates the body +// before it can consume capacity in the asynchronous queue. +func SanitizeOpsErrorBodyForQueue(raw string) (string, bool) { + return sanitizeErrorBodyForStorage(raw, OpsErrorLogQueueBodyMaxBytes) +} + +// SanitizeOpsUpstreamErrorsForQueue bounds and serializes attempt-level data +// before the entry can consume asynchronous queue capacity. +func SanitizeOpsUpstreamErrorsForQueue(entry *OpsInsertErrorLogInput) error { + return sanitizeOpsUpstreamErrors(entry) +} + func (s *OpsService) RecordError(ctx context.Context, entry *OpsInsertErrorLogInput) error { prepared, ok, err := s.prepareErrorLogInput(ctx, entry) if err != nil { @@ -295,7 +553,7 @@ func sanitizeOpsUpstreamErrors(entry *OpsInsertErrorLogInput) error { return nil } - const maxEvents = 32 + const maxEvents = 16 events := entry.UpstreamErrors if len(events) > maxEvents { events = events[len(events)-maxEvents:] @@ -308,9 +566,19 @@ func sanitizeOpsUpstreamErrors(entry *OpsInsertErrorLogInput) error { } out := *ev - out.Platform = strings.TrimSpace(out.Platform) + out.Platform = truncateString(strings.TrimSpace(out.Platform), 32) + out.AccountName = truncateString(strings.TrimSpace(out.AccountName), 128) out.UpstreamRequestID = truncateString(strings.TrimSpace(out.UpstreamRequestID), 128) + out.UpstreamURL = truncateString(strings.TrimSpace(out.UpstreamURL), 2048) + if body := strings.TrimSpace(out.UpstreamResponseBody); body != "" { + out.UpstreamResponseBody, _ = sanitizeErrorBodyForStorage(body, OpsErrorLogQueueBodyMaxBytes) + } else { + out.UpstreamResponseBody = "" + } out.Kind = truncateString(strings.TrimSpace(out.Kind), 64) + out.Stage = truncateString(strings.TrimSpace(out.Stage), 64) + out.Scope = truncateString(strings.TrimSpace(out.Scope), 64) + out.Reason = truncateString(strings.TrimSpace(out.Reason), 128) if out.AccountID < 0 { out.AccountID = 0 @@ -328,8 +596,8 @@ func sanitizeOpsUpstreamErrors(entry *OpsInsertErrorLogInput) error { detail := strings.TrimSpace(out.Detail) if detail != "" { - // Keep upstream detail small; request bodies are not stored here, only upstream error payloads. - sanitizedDetail, _ := sanitizeErrorBodyForStorage(detail, opsMaxStoredErrorBodyBytes) + // Keep upstream detail small while the event waits in the queue. + sanitizedDetail, _ := sanitizeErrorBodyForStorage(detail, OpsErrorLogQueueBodyMaxBytes) out.Detail = sanitizedDetail } else { out.Detail = "" @@ -375,8 +643,6 @@ func (s *OpsService) ListUserErrorRequests(ctx context.Context, userID int64, fi filter = &f uid := userID filter.UserID = &uid - // 用户侧放宽归属:纳入「删 key 后认证失败」(user_id=NULL,靠 deleted_key_owner 归因)的记录。 - filter.MatchDeletedKeyOwner = true // APIKeyID 透传:保留 handler 传入的值。安全由 buildOpsErrorLogsWhere 的 // "user_id = 自己 AND api_key_id = X" 双重约束保证——传入他人 key 只会得到空集,无泄露。 filter.View = "all" @@ -444,23 +710,14 @@ func (s *OpsService) GetUserErrorRequestDetail(ctx context.Context, userID, id i } return nil, infraerrors.InternalServer("OPS_ERROR_LOAD_FAILED", "Failed to load ops error log").WithCause(err) } - // 归属:直接归属(user_id)或经「已删除 key 归因」(deleted_key_owner_user_id)二者之一即可。 + // 归属只能由通过鉴权时写入的 user_id 确定。 ownedDirectly := detail.UserID != nil && *detail.UserID == userID - ownedViaDeletedKey := detail.DeletedKeyOwnerUserID != nil && *detail.DeletedKeyOwnerUserID == userID - if !ownedDirectly && !ownedViaDeletedKey { + if !ownedDirectly { return nil, infraerrors.NotFound("OPS_ERROR_NOT_FOUND", "ops error log not found") } return ToUserErrorRequestDetail(detail), nil } -// LookupDeletedKeyAudit 按明文 key 反查已删除 key 的原所有者;未命中或未启用返回 (nil, nil)。 -func (s *OpsService) LookupDeletedKeyAudit(ctx context.Context, key string) (*DeletedKeyAuditResult, error) { - if s.opsRepo == nil { - return nil, nil - } - return s.opsRepo.LookupDeletedKeyAudit(ctx, key) -} - func (s *OpsService) UpdateErrorResolution(ctx context.Context, errorID int64, resolved bool, resolvedByUserID *int64) error { if err := s.RequireMonitoringEnabled(ctx); err != nil { return err diff --git a/backend/internal/service/ops_service_user_error_test.go b/backend/internal/service/ops_service_user_error_test.go index c3b0967b6..b1f615270 100644 --- a/backend/internal/service/ops_service_user_error_test.go +++ b/backend/internal/service/ops_service_user_error_test.go @@ -162,61 +162,3 @@ func TestGetUserErrorRequestDetail_InvalidID(t *testing.T) { t.Fatal("expected error for id=-5") } } - -func TestListUserErrorRequests_EnablesMatchDeletedKeyOwner(t *testing.T) { - stub := &stubOpsRepoForUserErr{} - svc := &OpsService{opsRepo: stub} - uid := int64(42) - - if _, err := svc.ListUserErrorRequests(context.Background(), uid, &OpsErrorLogFilter{}); err != nil { - t.Fatal(err) - } - if stub.gotFilter == nil || !stub.gotFilter.MatchDeletedKeyOwner { - t.Fatal("ListUserErrorRequests should enable MatchDeletedKeyOwner for the user scope") - } -} - -func TestGetUserErrorRequestDetail_DeletedKeyOwnerAccess(t *testing.T) { - ownerUID := int64(777) - otherUID := int64(2) - - // 情况2:user_id=NULL,靠 deleted_key_owner_user_id 归因到 ownerUID - mk := func() *OpsErrorLogDetail { - return &OpsErrorLogDetail{ - OpsErrorLog: OpsErrorLog{ - ID: 55, - Phase: "auth", - Type: "api_error", - StatusCode: 401, - Message: "Invalid API key", - UserID: nil, - APIKeyName: "my-old-key", - APIKeyDeleted: true, - DeletedKeyOwnerUserID: &ownerUID, - }, - } - } - - // 原所有者(经 deleted_key 归因)→ 放行 - svcOwner := &OpsService{opsRepo: &stubOpsRepoForUserErr{detailToReturn: mk()}} - got, err := svcOwner.GetUserErrorRequestDetail(context.Background(), ownerUID, 55) - if err != nil { - t.Fatalf("owner via deleted_key should be allowed, got err: %v", err) - } - if got == nil || got.ID != 55 { - t.Fatalf("expected detail ID=55, got %+v", got) - } - if !got.KeyDeleted || got.KeyName != "my-old-key" { - t.Fatalf("expected KeyDeleted=true KeyName=my-old-key, got %+v", got) - } - - // 他人 → NotFound,不泄露存在性 - svcOther := &OpsService{opsRepo: &stubOpsRepoForUserErr{detailToReturn: mk()}} - got2, err2 := svcOther.GetUserErrorRequestDetail(context.Background(), otherUID, 55) - if err2 == nil || got2 != nil { - t.Fatalf("non-owner should get (nil, NotFound), got detail=%+v err=%v", got2, err2) - } - if !infraerrors.IsNotFound(err2) { - t.Fatalf("expected NotFound, got %v", err2) - } -} diff --git a/backend/internal/service/ops_settings.go b/backend/internal/service/ops_settings.go index 472f4e320..e2943be5c 100644 --- a/backend/internal/service/ops_settings.go +++ b/backend/internal/service/ops_settings.go @@ -373,6 +373,7 @@ func defaultOpsAdvancedSettings() *OpsAdvancedSettings { IgnoreCountTokensErrors: true, // count_tokens 404 是预期行为,默认忽略 IgnoreContextCanceled: true, // Default to true - client disconnects are not errors IgnoreNoAvailableAccounts: false, // Default to false - this is a real routing issue + IgnoreInvalidApiKeyErrors: true, // Legacy compatibility field; admission rejects are always excluded. IgnoreInsufficientBalanceErrors: false, // 默认不忽略,余额不足可能需要关注 DisplayOpenAITokenStats: false, DisplayAlertEvents: true, @@ -385,6 +386,10 @@ func normalizeOpsAdvancedSettings(cfg *OpsAdvancedSettings) { if cfg == nil { return } + // Admission rejects are a security/traffic concern, not an operational + // request-error category. Keep the legacy field true for old clients but do + // not allow it to re-enable those rows. + cfg.IgnoreInvalidApiKeyErrors = true cfg.OpenAIAccountQuotaAutoPause.DefaultThreshold5h = clampOpsQuotaAutoPauseThreshold(cfg.OpenAIAccountQuotaAutoPause.DefaultThreshold5h) cfg.OpenAIAccountQuotaAutoPause.DefaultThreshold7d = clampOpsQuotaAutoPauseThreshold(cfg.OpenAIAccountQuotaAutoPause.DefaultThreshold7d) cfg.DataRetention.CleanupSchedule = strings.TrimSpace(cfg.DataRetention.CleanupSchedule) @@ -439,32 +444,20 @@ func validateOpsAdvancedSettings(cfg *OpsAdvancedSettings) error { } func (s *OpsService) GetOpsAdvancedSettings(ctx context.Context) (*OpsAdvancedSettings, error) { - defaultCfg := defaultOpsAdvancedSettings() - if s == nil || s.settingRepo == nil { - return defaultCfg, nil - } - if ctx == nil { - ctx = context.Background() - } + _ = ctx + cfg := s.OpsAdvancedSettingsSnapshot() + return &cfg, nil +} - raw, err := s.settingRepo.GetValue(ctx, SettingKeyOpsAdvancedSettings) - if err != nil { - if errors.Is(err, ErrSettingNotFound) { - if b, mErr := json.Marshal(defaultCfg); mErr == nil { - _ = s.settingRepo.Set(ctx, SettingKeyOpsAdvancedSettings, string(b)) - } - return defaultCfg, nil +// OpsAdvancedSettingsSnapshot returns a value copy for request hot paths. It +// avoids both repository I/O and pointer escape/allocation. +func (s *OpsService) OpsAdvancedSettingsSnapshot() OpsAdvancedSettings { + if s != nil { + if snapshot := s.runtimeSettings.Load(); snapshot != nil { + return snapshot.advanced } - return nil, err } - - cfg := defaultOpsAdvancedSettings() - if err := json.Unmarshal([]byte(raw), cfg); err != nil { - return defaultCfg, nil - } - - normalizeOpsAdvancedSettings(cfg) - return cfg, nil + return *defaultOpsAdvancedSettings() } func (s *OpsService) UpdateOpsAdvancedSettings(ctx context.Context, cfg *OpsAdvancedSettings) (*OpsAdvancedSettings, error) { @@ -490,6 +483,7 @@ func (s *OpsService) UpdateOpsAdvancedSettings(ctx context.Context, cfg *OpsAdva if err := s.settingRepo.Set(ctx, SettingKeyOpsAdvancedSettings, string(raw)); err != nil { return nil, err } + s.storeAdvancedSettingsSnapshot(cfg) // Push the new quota auto-pause settings straight into the in-memory cache that // the OpenAI scheduling hot path reads, so the next request observes the new value // without waiting for the background refresher's TTL. diff --git a/backend/internal/service/ops_settings_advanced_test.go b/backend/internal/service/ops_settings_advanced_test.go index 62803f948..39d0b8722 100644 --- a/backend/internal/service/ops_settings_advanced_test.go +++ b/backend/internal/service/ops_settings_advanced_test.go @@ -9,7 +9,7 @@ import ( "github.com/Wei-Shaw/sub2api/internal/config" ) -func TestGetOpsAdvancedSettings_DefaultHidesOpenAITokenStats(t *testing.T) { +func TestGetOpsAdvancedSettings_DefaultSnapshotHidesOpenAITokenStats(t *testing.T) { repo := newRuntimeSettingRepoStub() svc := &OpsService{settingRepo: repo} @@ -23,8 +23,8 @@ func TestGetOpsAdvancedSettings_DefaultHidesOpenAITokenStats(t *testing.T) { if !cfg.DisplayAlertEvents { t.Fatalf("DisplayAlertEvents = false, want true by default") } - if repo.setCalls != 1 { - t.Fatalf("expected defaults to be persisted once, got %d", repo.setCalls) + if repo.getValueCalls != 0 || repo.getMultipleCalls != 0 { + t.Fatalf("hot-path snapshot read touched repository: get=%d get_multiple=%d", repo.getValueCalls, repo.getMultipleCalls) } } @@ -46,6 +46,7 @@ func TestUpdateOpsAdvancedSettings_PersistsOpenAITokenStatsVisibility(t *testing if updated.DisplayAlertEvents { t.Fatalf("DisplayAlertEvents = true, want false") } + readsAfterUpdate := repo.getValueCalls + repo.getMultipleCalls reloaded, err := svc.GetOpsAdvancedSettings(context.Background()) if err != nil { @@ -57,6 +58,9 @@ func TestUpdateOpsAdvancedSettings_PersistsOpenAITokenStatsVisibility(t *testing if reloaded.DisplayAlertEvents { t.Fatalf("reloaded DisplayAlertEvents = true, want false") } + if got := repo.getValueCalls + repo.getMultipleCalls; got != readsAfterUpdate { + t.Fatalf("snapshot reload performed repository read: before=%d after=%d", readsAfterUpdate, got) + } } func TestGetOpsAdvancedSettings_BackfillsNewDisplayFlagsFromDefaults(t *testing.T) { @@ -77,7 +81,7 @@ func TestGetOpsAdvancedSettings_BackfillsNewDisplayFlagsFromDefaults(t *testing. "ignore_count_tokens_errors": true, "ignore_context_canceled": true, "ignore_no_available_accounts": false, - "ignore_invalid_api_key_errors": false, + "ignore_invalid_api_key_errors": true, "auto_refresh_enabled": false, "auto_refresh_interval_seconds": 30, } diff --git a/backend/internal/service/ops_settings_models.go b/backend/internal/service/ops_settings_models.go index 4d459e2b9..305797800 100644 --- a/backend/internal/service/ops_settings_models.go +++ b/backend/internal/service/ops_settings_models.go @@ -92,18 +92,19 @@ type OpsAlertRuntimeSettings struct { // OpsAdvancedSettings stores advanced ops configuration (data retention, aggregation). type OpsAdvancedSettings struct { - DataRetention OpsDataRetentionSettings `json:"data_retention"` - Aggregation OpsAggregationSettings `json:"aggregation"` - OpenAIAccountQuotaAutoPause OpsOpenAIAccountQuotaAutoPauseSettings `json:"openai_account_quota_auto_pause"` - IgnoreCountTokensErrors bool `json:"ignore_count_tokens_errors"` - IgnoreContextCanceled bool `json:"ignore_context_canceled"` - IgnoreNoAvailableAccounts bool `json:"ignore_no_available_accounts"` - IgnoreInvalidApiKeyErrors bool `json:"ignore_invalid_api_key_errors"` - IgnoreInsufficientBalanceErrors bool `json:"ignore_insufficient_balance_errors"` - DisplayOpenAITokenStats bool `json:"display_openai_token_stats"` - DisplayAlertEvents bool `json:"display_alert_events"` - AutoRefreshEnabled bool `json:"auto_refresh_enabled"` - AutoRefreshIntervalSec int `json:"auto_refresh_interval_seconds"` + DataRetention OpsDataRetentionSettings `json:"data_retention"` + Aggregation OpsAggregationSettings `json:"aggregation"` + OpenAIAccountQuotaAutoPause OpsOpenAIAccountQuotaAutoPauseSettings `json:"openai_account_quota_auto_pause"` + IgnoreCountTokensErrors bool `json:"ignore_count_tokens_errors"` + IgnoreContextCanceled bool `json:"ignore_context_canceled"` + IgnoreNoAvailableAccounts bool `json:"ignore_no_available_accounts"` + // Deprecated compatibility field. It is always normalized to true. + IgnoreInvalidApiKeyErrors bool `json:"ignore_invalid_api_key_errors"` + IgnoreInsufficientBalanceErrors bool `json:"ignore_insufficient_balance_errors"` + DisplayOpenAITokenStats bool `json:"display_openai_token_stats"` + DisplayAlertEvents bool `json:"display_alert_events"` + AutoRefreshEnabled bool `json:"auto_refresh_enabled"` + AutoRefreshIntervalSec int `json:"auto_refresh_interval_seconds"` } type OpsOpenAIAccountQuotaAutoPauseSettings struct { diff --git a/backend/internal/service/ops_system_log_sink.go b/backend/internal/service/ops_system_log_sink.go index 2e6f5515c..1cbac237c 100644 --- a/backend/internal/service/ops_system_log_sink.go +++ b/backend/internal/service/ops_system_log_sink.go @@ -112,6 +112,11 @@ func (s *OpsSystemLogSink) WriteLogEvent(event *logger.LogEvent) { } func (s *OpsSystemLogSink) shouldIndex(event *logger.LogEvent) bool { + if event != nil && event.Fields != nil { + if skip, _ := event.Fields[logger.OpsSystemLogSkipField].(bool); skip { + return false + } + } level := strings.ToLower(strings.TrimSpace(event.Level)) switch level { case "warn", "warning", "error", "fatal", "panic", "dpanic": diff --git a/backend/internal/service/ops_system_log_sink_test.go b/backend/internal/service/ops_system_log_sink_test.go index 0d15f1a66..fbe61d386 100644 --- a/backend/internal/service/ops_system_log_sink_test.go +++ b/backend/internal/service/ops_system_log_sink_test.go @@ -36,6 +36,15 @@ func TestOpsSystemLogSink_ShouldIndex(t *testing.T) { event: &logger.LogEvent{Level: "info", Component: "http.access"}, want: true, }, + { + name: "rejected access excluded from database sink", + event: &logger.LogEvent{ + Level: "info", + Component: "http.access", + Fields: map[string]any{logger.OpsSystemLogSkipField: true}, + }, + want: false, + }, { name: "access component from fields (real zap path)", event: &logger.LogEvent{ diff --git a/backend/internal/service/ops_upstream_context.go b/backend/internal/service/ops_upstream_context.go index 2d28f1a16..35a97cc2f 100644 --- a/backend/internal/service/ops_upstream_context.go +++ b/backend/internal/service/ops_upstream_context.go @@ -102,28 +102,54 @@ func HasOpsClientBusinessLimited(c *gin.Context) bool { type OpsStreamError struct { // ErrType 是写入 SSE 帧的对客错误类型(如 rate_limit_error / upstream_error / api_error)。 ErrType string + // Code 是可选的稳定错误分类;用于既保留通用 OpenAI error.type,又向客户端和 Ops + // 暴露可编程判断的细分类(如 upstream_http2_stream_error)。 + Code string // Message 是写入 SSE 帧的对客错误消息。 Message string // IntendedStatus 是流若未固化本应返回的 HTTP 状态码(如并发限流的 429)。 - // 仅用于错误分级(severity/classification);实际 wire 状态码仍为 200。 + // 默认仅用于错误分级;CountTowardsSLA=true 时也作为 Ops 的逻辑状态码。 IntendedStatus int + // CountTowardsSLA 表示虽然 wire 状态已固化为 200,请求在应用语义上仍然失败, + // Ops 应使用 IntendedStatus 计入错误率/SLA。 + CountTowardsSLA bool } // MarkOpsStreamError 记录一次就地 SSE 错误,供 ops 日志采集。 // 采用「首个标记生效」策略:同一请求若先后补发多帧(如上游透传错误后又追加通用兜底帧), // 保留最先记录的根因错误,而不是被后续的 "Upstream request failed" 覆盖。 func MarkOpsStreamError(c *gin.Context, errType, message string, intendedStatus int) { + markOpsStreamError(c, OpsStreamError{ + ErrType: errType, + Message: message, + IntendedStatus: intendedStatus, + }) +} + +// MarkOpsStreamFailure records an in-band stream error that represents a failed +// request and therefore must count towards Ops error rate/SLA despite HTTP 200 +// already being committed on the wire. +func MarkOpsStreamFailure(c *gin.Context, errType, code, message string, intendedStatus int) { + markOpsStreamError(c, OpsStreamError{ + ErrType: errType, + Code: code, + Message: message, + IntendedStatus: intendedStatus, + CountTowardsSLA: true, + }) +} + +func markOpsStreamError(c *gin.Context, streamErr OpsStreamError) { if c == nil { return } if _, exists := c.Get(OpsStreamErrorKey); exists { return } - c.Set(OpsStreamErrorKey, OpsStreamError{ - ErrType: strings.TrimSpace(errType), - Message: strings.TrimSpace(message), - IntendedStatus: intendedStatus, - }) + streamErr.ErrType = strings.TrimSpace(streamErr.ErrType) + streamErr.Code = strings.TrimSpace(streamErr.Code) + streamErr.Message = strings.TrimSpace(streamErr.Message) + c.Set(OpsStreamErrorKey, streamErr) } // GetOpsStreamError 返回本请求记录的就地 SSE 错误(若有)。 diff --git a/backend/internal/service/ratelimit_service.go b/backend/internal/service/ratelimit_service.go index 507e02d87..266f806e8 100644 --- a/backend/internal/service/ratelimit_service.go +++ b/backend/internal/service/ratelimit_service.go @@ -150,7 +150,8 @@ const ( // CheckErrorPolicy 检查自定义错误码和临时不可调度规则。 // 自定义错误码开启时覆盖后续所有逻辑(包括临时不可调度)。 -func (s *RateLimitService) CheckErrorPolicy(ctx context.Context, account *Account, statusCode int, responseBody []byte) ErrorPolicyResult { +func (s *RateLimitService) CheckErrorPolicy(ctx context.Context, account *Account, statusCode int, responseBody []byte, requestedModel ...string) ErrorPolicyResult { + ctx = withTempUnschedulableModel(ctx, requestedModel) if account.IsCustomErrorCodesEnabled() { if account.ShouldHandleErrorCode(statusCode) { return ErrorPolicyMatched @@ -159,9 +160,14 @@ func (s *RateLimitService) CheckErrorPolicy(ctx context.Context, account *Accoun return ErrorPolicySkipped } if account.IsPoolMode() { + // 池模式只跳过默认账号状态处理;管理员显式配置的临时不可调度规则仍应生效。 + // 401 保留现有认证错误语义,避免改变重复 401 的升级行为。 + if statusCode != http.StatusUnauthorized && s.tryTempUnschedulable(ctx, account, statusCode, responseBody) { + return ErrorPolicyTempUnscheduled + } return ErrorPolicySkipped } - if s.tryTempUnschedulable(ctx, account, statusCode, responseBody) { + if s.tryTempUnschedulable(ctx, account, statusCode, responseBody, firstRequestedModel(requestedModel)) { return ErrorPolicyTempUnscheduled } return ErrorPolicyNone @@ -170,10 +176,15 @@ func (s *RateLimitService) CheckErrorPolicy(ctx context.Context, account *Accoun // HandleUpstreamError 处理上游错误响应,标记账号状态 // 返回是否应该停止该账号的调度 func (s *RateLimitService) HandleUpstreamError(ctx context.Context, account *Account, statusCode int, headers http.Header, responseBody []byte, requestedModel ...string) (shouldDisable bool) { + ctx = withTempUnschedulableModel(ctx, requestedModel) customErrorCodesEnabled := account.IsCustomErrorCodesEnabled() - // 池模式默认不标记本地账号状态;仅当用户显式配置自定义错误码时按本地策略处理。 + // 池模式默认不标记本地账号状态;但管理员显式配置的临时不可调度规则优先。 + // 401 保留现有认证错误语义,不在这里改变池模式的认证处理。 if account.IsPoolMode() && !customErrorCodesEnabled { + if statusCode != http.StatusUnauthorized && s.tryTempUnschedulable(ctx, account, statusCode, responseBody) { + return true + } slog.Info("pool_mode_error_skipped", "account_id", account.ID, "status_code", statusCode) return false } @@ -207,7 +218,7 @@ func (s *RateLimitService) HandleUpstreamError(ctx context.Context, account *Acc // 先尝试临时不可调度规则(401除外) // 如果匹配成功,直接返回,不执行后续禁用逻辑 if statusCode != 401 { - if s.tryTempUnschedulable(ctx, account, statusCode, responseBody) { + if s.tryTempUnschedulable(ctx, account, statusCode, responseBody, firstRequestedModel(requestedModel)) { return true } } @@ -1891,14 +1902,18 @@ func (s *RateLimitService) GetTempUnschedStatus(ctx context.Context, accountID i return state, nil } -func (s *RateLimitService) HandleTempUnschedulable(ctx context.Context, account *Account, statusCode int, responseBody []byte) bool { +func (s *RateLimitService) HandleTempUnschedulable(ctx context.Context, account *Account, statusCode int, responseBody []byte, requestedModel ...string) bool { if account == nil { return false } + if account.IsPoolMode() && !account.IsCustomErrorCodesEnabled() { + return false + } if !account.ShouldHandleErrorCode(statusCode) { return false } - return s.tryTempUnschedulable(ctx, account, statusCode, responseBody) + ctx = withTempUnschedulableModel(ctx, requestedModel) + return s.tryTempUnschedulable(ctx, account, statusCode, responseBody, firstRequestedModel(requestedModel)) } func (s *RateLimitService) HandleOpenAIImageRateLimit(ctx context.Context, account *Account, statusCode int, headers http.Header, responseBody []byte) bool { @@ -2066,7 +2081,71 @@ func modelRateLimitKeyForUpstreamModelNotFound(ctx context.Context, account *Acc return modelKey } -func (s *RateLimitService) tryTempUnschedulable(ctx context.Context, account *Account, statusCode int, responseBody []byte) bool { +func firstRequestedModel(requestedModel []string) string { + if len(requestedModel) == 0 { + return "" + } + return strings.TrimSpace(requestedModel[0]) +} + +type tempUnschedulableModelContextKey struct{} + +func withTempUnschedulableModel(ctx context.Context, requestedModel []string) context.Context { + model := firstRequestedModel(requestedModel) + if model == "" { + return ctx + } + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, tempUnschedulableModelContextKey{}, model) +} + +func tempUnschedulableModel(ctx context.Context, requestedModel []string) string { + if model := firstRequestedModel(requestedModel); model != "" { + return model + } + if ctx == nil { + return "" + } + model, _ := ctx.Value(tempUnschedulableModelContextKey{}).(string) + return strings.TrimSpace(model) +} + +type tempUnschedulableRuleMatch struct { + rule TempUnschedulableRule + ruleIndex int + matchedKeyword string +} + +func matchTempUnschedulableRules(account *Account, statusCode int, responseBody []byte) []tempUnschedulableRuleMatch { + if account == nil || !account.IsTempUnschedulableEnabled() || statusCode <= 0 || len(responseBody) == 0 { + return nil + } + rules := account.GetTempUnschedulableRules() + if len(rules) == 0 { + return nil + } + body := responseBody + if len(body) > tempUnschedBodyMaxBytes { + body = body[:tempUnschedBodyMaxBytes] + } + bodyLower := strings.ToLower(string(body)) + matches := make([]tempUnschedulableRuleMatch, 0, 1) + for idx, rule := range rules { + if rule.ErrorCode != statusCode || len(rule.Keywords) == 0 { + continue + } + matchedKeyword := matchTempUnschedKeyword(bodyLower, rule.Keywords) + if matchedKeyword == "" { + continue + } + matches = append(matches, tempUnschedulableRuleMatch{rule: rule, ruleIndex: idx, matchedKeyword: matchedKeyword}) + } + return matches +} + +func (s *RateLimitService) tryTempUnschedulable(ctx context.Context, account *Account, statusCode int, responseBody []byte, requestedModel ...string) bool { if account == nil { return false } @@ -2090,30 +2169,8 @@ func (s *RateLimitService) tryTempUnschedulable(ctx context.Context, account *Ac return false } } - rules := account.GetTempUnschedulableRules() - if len(rules) == 0 { - return false - } - if statusCode <= 0 || len(responseBody) == 0 { - return false - } - - body := responseBody - if len(body) > tempUnschedBodyMaxBytes { - body = body[:tempUnschedBodyMaxBytes] - } - bodyLower := strings.ToLower(string(body)) - - for idx, rule := range rules { - if rule.ErrorCode != statusCode || len(rule.Keywords) == 0 { - continue - } - matchedKeyword := matchTempUnschedKeyword(bodyLower, rule.Keywords) - if matchedKeyword == "" { - continue - } - - if s.triggerTempUnschedulable(ctx, account, rule, idx, statusCode, matchedKeyword, responseBody) { + for _, match := range matchTempUnschedulableRules(account, statusCode, responseBody) { + if s.triggerTempUnschedulable(ctx, account, match.rule, match.ruleIndex, statusCode, match.matchedKeyword, responseBody, tempUnschedulableModel(ctx, requestedModel)) { return true } } @@ -2153,7 +2210,7 @@ func matchTempUnschedKeyword(bodyLower string, keywords []string) string { return "" } -func (s *RateLimitService) triggerTempUnschedulable(ctx context.Context, account *Account, rule TempUnschedulableRule, ruleIndex int, statusCode int, matchedKeyword string, responseBody []byte) bool { +func (s *RateLimitService) triggerTempUnschedulable(ctx context.Context, account *Account, rule TempUnschedulableRule, ruleIndex int, statusCode int, matchedKeyword string, responseBody []byte, requestedModel ...string) bool { if account == nil { return false } @@ -2181,6 +2238,21 @@ func (s *RateLimitService) triggerTempUnschedulable(ctx context.Context, account reason = strings.TrimSpace(state.ErrorMessage) } + // Persist known-model failures under the model key so the scheduler excludes + // only this (account, model) pair. Authentication and model-unknown failures + // retain the legacy account-wide temporary-unschedulable behavior below. + modelKey := firstRequestedModel(requestedModel) + if modelKey != "" && statusCode != http.StatusUnauthorized { + if err := s.accountRepo.SetModelRateLimit(ctx, account.ID, modelKey, until, reason); err != nil { + slog.Warn("temp_unsched_model_rate_limit_set_failed", "account_id", account.ID, "model", modelKey, "error", err) + // The rule matched, so fail over the current request even if persistence + // failed; never widen a model-scoped failure into an account-wide block. + return true + } + slog.Info("account_model_temp_unschedulable", "account_id", account.ID, "model", modelKey, "until", until, "rule_index", ruleIndex, "status_code", statusCode) + return true + } + s.notifyAccountSchedulingBlocked(account, until, "temp_unschedulable") if err := s.accountRepo.SetTempUnschedulable(ctx, account.ID, until, reason); err != nil { slog.Warn("temp_unsched_set_failed", "account_id", account.ID, "error", err) diff --git a/backend/internal/service/ratelimit_service_model_not_found_test.go b/backend/internal/service/ratelimit_service_model_not_found_test.go index 51bd8a607..55becb28a 100644 --- a/backend/internal/service/ratelimit_service_model_not_found_test.go +++ b/backend/internal/service/ratelimit_service_model_not_found_test.go @@ -4,6 +4,7 @@ package service import ( "context" + "encoding/json" "errors" "net/http" "testing" @@ -87,7 +88,7 @@ func TestRateLimitService_HandleUpstreamError_ModelNotFoundWriteFailureDoesNotTe require.Len(t, repo.modelRateLimitCalls, 1) } -func TestRateLimitService_HandleUpstreamError_Bare404KeepsTempUnschedulablePath(t *testing.T) { +func TestRateLimitService_HandleUpstreamError_Bare404UsesModelScopedTempUnschedulableWhenModelKnown(t *testing.T) { repo := &modelNotFoundAccountRepoStub{} svc := &RateLimitService{accountRepo: repo} account := openAIModelNotFoundTempAccount() @@ -101,11 +102,203 @@ func TestRateLimitService_HandleUpstreamError_Bare404KeepsTempUnschedulablePath( "gpt-5.4", ) + require.True(t, handled) + require.Zero(t, repo.tempCalls) + require.Len(t, repo.modelRateLimitCalls, 1) + call := repo.modelRateLimitCalls[0] + require.Equal(t, account.ID, call.accountID) + require.Equal(t, "gpt-5.4", call.scope) + require.WithinDuration(t, time.Now().Add(10*time.Minute), call.resetAt, 5*time.Second) + + var state TempUnschedState + require.NoError(t, json.Unmarshal([]byte(call.reason), &state)) + require.Equal(t, http.StatusNotFound, state.StatusCode) + require.Equal(t, "not found", state.MatchedKeyword) +} + +func TestRateLimitService_HandleUpstreamError_Bare404WithoutModelKeepsAccountTempUnschedulable(t *testing.T) { + repo := &modelNotFoundAccountRepoStub{} + svc := &RateLimitService{accountRepo: repo} + account := openAIModelNotFoundTempAccount() + + handled := svc.HandleUpstreamError( + context.Background(), + account, + http.StatusNotFound, + http.Header{}, + []byte(`{"error":{"message":"endpoint not found"}}`), + ) + require.True(t, handled) require.Equal(t, 1, repo.tempCalls) require.Empty(t, repo.modelRateLimitCalls) } +func TestRateLimitService_HandleUpstreamError_ModelTempWriteFailureNeverWidensToAccount(t *testing.T) { + repo := &modelNotFoundAccountRepoStub{modelRateLimitErr: errors.New("write failed")} + svc := &RateLimitService{accountRepo: repo} + account := openAIModelNotFoundTempAccount() + + handled := svc.HandleUpstreamError( + context.Background(), + account, + http.StatusNotFound, + http.Header{}, + []byte(`{"error":{"message":"endpoint not found"}}`), + "gpt-5.4", + ) + + require.True(t, handled) + require.Zero(t, repo.tempCalls) + require.Len(t, repo.modelRateLimitCalls, 1) +} + +func TestRateLimitService_HandleTempUnschedulable_PoolModeWithoutCustomPolicySkipsState(t *testing.T) { + repo := &modelNotFoundAccountRepoStub{} + svc := &RateLimitService{accountRepo: repo} + account := openAIModelNotFoundTempAccount() + account.Credentials["pool_mode"] = true + + handled := svc.HandleTempUnschedulable( + context.Background(), + account, + http.StatusNotFound, + []byte(`{"error":{"message":"endpoint not found"}}`), + "gpt-5.4", + ) + + require.False(t, handled) + require.Zero(t, repo.tempCalls) + require.Empty(t, repo.modelRateLimitCalls) +} + +func TestRateLimitService_HandleTempUnschedulable_PoolModeCustomPolicyUsesModelScope(t *testing.T) { + repo := &modelNotFoundAccountRepoStub{} + svc := &RateLimitService{accountRepo: repo} + account := openAIModelNotFoundTempAccount() + account.Credentials["pool_mode"] = true + account.Credentials["custom_error_codes_enabled"] = true + account.Credentials["custom_error_codes"] = []any{float64(http.StatusNotFound)} + + handled := svc.HandleTempUnschedulable( + context.Background(), + account, + http.StatusNotFound, + []byte(`{"error":{"message":"endpoint not found"}}`), + "gpt-5.4", + ) + + require.True(t, handled) + require.Zero(t, repo.tempCalls) + require.Len(t, repo.modelRateLimitCalls, 1) + require.Equal(t, "gpt-5.4", repo.modelRateLimitCalls[0].scope) +} + +func TestRateLimitService_TempUnschedulableContextPreservesModelForPoolDependency(t *testing.T) { + repo := &modelNotFoundAccountRepoStub{} + svc := &RateLimitService{accountRepo: repo} + account := openAIModelNotFoundTempAccount() + account.Credentials["pool_mode"] = true + ctx := withTempUnschedulableModel(context.Background(), []string{"gpt-5.4"}) + + // #4496 calls tryTempUnschedulable from the pool-mode branch without an + // explicit model argument. The request context must preserve the canonical + // model so that combined behavior remains model-scoped after it lands. + handled := svc.tryTempUnschedulable( + ctx, + account, + http.StatusNotFound, + []byte(`{"error":{"message":"endpoint not found"}}`), + ) + + require.True(t, handled) + require.Zero(t, repo.tempCalls) + require.Len(t, repo.modelRateLimitCalls, 1) + require.Equal(t, "gpt-5.4", repo.modelRateLimitCalls[0].scope) +} + +func TestRateLimitService_HandleUpstreamError_CustomPolicyExclusionSkipsAllState(t *testing.T) { + repo := &modelNotFoundAccountRepoStub{} + svc := &RateLimitService{accountRepo: repo} + account := openAIModelNotFoundTempAccount() + account.Credentials["custom_error_codes_enabled"] = true + account.Credentials["custom_error_codes"] = []any{float64(http.StatusServiceUnavailable)} + + handled := svc.HandleUpstreamError( + context.Background(), + account, + http.StatusNotFound, + http.Header{}, + []byte(`{"error":{"message":"endpoint not found"}}`), + "gpt-5.4", + ) + + require.False(t, handled) + require.Zero(t, repo.tempCalls) + require.Empty(t, repo.modelRateLimitCalls) +} + +func TestRateLimitService_HandleTempUnschedulable_AuthenticationFailureStaysAccountScoped(t *testing.T) { + repo := &modelNotFoundAccountRepoStub{} + svc := &RateLimitService{accountRepo: repo} + account := openAIModelNotFoundTempAccount() + account.TempUnschedulableReason = "legacy non-JSON reason" + account.Credentials["temp_unschedulable_rules"] = []any{ + map[string]any{ + "error_code": float64(http.StatusUnauthorized), + "keywords": []any{"unauthorized"}, + "duration_minutes": float64(10), + }, + } + + handled := svc.HandleTempUnschedulable( + context.Background(), + account, + http.StatusUnauthorized, + []byte(`{"error":{"message":"unauthorized"}}`), + "gpt-5.4", + ) + + require.True(t, handled) + require.Equal(t, 1, repo.tempCalls) + require.Empty(t, repo.modelRateLimitCalls) +} + +func TestRateLimitService_ModelTempUnschedulableIsolatesSchedulerByModel(t *testing.T) { + repo := &modelNotFoundAccountRepoStub{} + svc := &RateLimitService{accountRepo: repo} + account := openAIModelNotFoundTempAccount() + account.Credentials["model_mapping"] = map[string]any{ + "public-a": "upstream-a", + "upstream-a": "upstream-b", + } + + handled := svc.HandleUpstreamError( + context.Background(), + account, + http.StatusNotFound, + http.Header{}, + []byte(`{"error":{"message":"endpoint not found"}}`), + "upstream-a", + ) + + require.True(t, handled) + require.Len(t, repo.modelRateLimitCalls, 1) + call := repo.modelRateLimitCalls[0] + require.Equal(t, "upstream-a", call.scope, "canonical upstream model must not be mapped a second time") + + account.Extra = map[string]any{ + modelRateLimitsKey: map[string]any{ + call.scope: map[string]any{ + "rate_limit_reset_at": call.resetAt.UTC().Format(time.RFC3339), + }, + }, + } + + require.False(t, account.IsSchedulableForModelWithContext(context.Background(), "public-a")) + require.True(t, account.IsSchedulableForModelWithContext(context.Background(), "gpt-5.6-sol")) +} + func openAIModelNotFoundTempAccount() *Account { return &Account{ ID: 101, diff --git a/backend/internal/service/ratelimit_session_window_test.go b/backend/internal/service/ratelimit_session_window_test.go index 279a31ccd..b48412114 100644 --- a/backend/internal/service/ratelimit_session_window_test.go +++ b/backend/internal/service/ratelimit_session_window_test.go @@ -143,6 +143,9 @@ func (m *sessionWindowMockRepo) ListSchedulableUngroupedByPlatform(context.Conte func (m *sessionWindowMockRepo) ListSchedulableUngroupedByPlatforms(context.Context, []string) ([]Account, error) { panic("unexpected") } +func (m *sessionWindowMockRepo) ListModelAvailabilityCandidates(context.Context, *int64, []string, bool) ([]Account, error) { + panic("unexpected") +} func (m *sessionWindowMockRepo) SetRateLimited(context.Context, int64, time.Time) error { panic("unexpected") } diff --git a/backend/internal/service/scheduler_snapshot_batch_query_test.go b/backend/internal/service/scheduler_snapshot_batch_query_test.go index 221610bdb..b957891d3 100644 --- a/backend/internal/service/scheduler_snapshot_batch_query_test.go +++ b/backend/internal/service/scheduler_snapshot_batch_query_test.go @@ -56,6 +56,10 @@ func (r *batchAccountQueryRepo) ListSchedulableUngroupedByPlatforms(_ context.Co return r.run(batchAccountQueryKey{platform: platforms[0], mixed: true}) } +func (r *batchAccountQueryRepo) ListModelAvailabilityCandidates(context.Context, *int64, []string, bool) ([]Account, error) { + panic("unexpected ListModelAvailabilityCandidates call") +} + func (r *batchAccountQueryRepo) ListSchedulableByPlatform(_ context.Context, platform string) ([]Account, error) { return r.run(batchAccountQueryKey{platform: platform}) } diff --git a/backend/internal/service/scheduler_snapshot_full_rebuild_lifecycle_test.go b/backend/internal/service/scheduler_snapshot_full_rebuild_lifecycle_test.go index 35ce99a8d..fa918a086 100644 --- a/backend/internal/service/scheduler_snapshot_full_rebuild_lifecycle_test.go +++ b/backend/internal/service/scheduler_snapshot_full_rebuild_lifecycle_test.go @@ -214,6 +214,10 @@ func (r *fullRebuildAccountRepo) ListSchedulableUngroupedByPlatforms(_ context.C return r.record(0, firstPlatform(platforms)) } +func (r *fullRebuildAccountRepo) ListModelAvailabilityCandidates(context.Context, *int64, []string, bool) ([]Account, error) { + panic("unexpected ListModelAvailabilityCandidates call") +} + func (r *fullRebuildAccountRepo) ListSchedulableByPlatform(_ context.Context, platform string) ([]Account, error) { return r.record(0, platform) } diff --git a/backend/internal/service/subscription_assign_idempotency_test.go b/backend/internal/service/subscription_assign_idempotency_test.go index d4913f899..151b5c792 100644 --- a/backend/internal/service/subscription_assign_idempotency_test.go +++ b/backend/internal/service/subscription_assign_idempotency_test.go @@ -268,7 +268,7 @@ func (s *subscriptionUserSubRepoStub) Update(_ context.Context, sub *UserSubscri } func TestAssignSubscriptionReuseWhenSemanticsMatch(t *testing.T) { - start := time.Date(2026, 2, 20, 10, 0, 0, 0, time.UTC) + start := time.Now().Add(-time.Hour) groupRepo := &subscriptionGroupRepoStub{ group: &Group{ID: 1, SubscriptionType: SubscriptionTypeSubscription}, } @@ -279,6 +279,7 @@ func TestAssignSubscriptionReuseWhenSemanticsMatch(t *testing.T) { GroupID: 1, StartsAt: start, ExpiresAt: start.AddDate(0, 0, 30), + Status: SubscriptionStatusActive, Notes: "init", }) @@ -292,10 +293,169 @@ func TestAssignSubscriptionReuseWhenSemanticsMatch(t *testing.T) { require.NoError(t, err) require.Equal(t, int64(10), sub.ID) require.Equal(t, 0, subRepo.createCalls, "reuse should not create new subscription") + require.Equal(t, start, sub.StartsAt) + require.Equal(t, start.AddDate(0, 0, 30), sub.ExpiresAt) +} + +func TestAssignSubscriptionDoesNotReactivateFutureSuspendedSubscription(t *testing.T) { + start := time.Now().Add(-time.Hour) + groupRepo := &subscriptionGroupRepoStub{ + group: &Group{ID: 1, SubscriptionType: SubscriptionTypeSubscription}, + } + subRepo := newSubscriptionUserSubRepoStub() + subRepo.seed(&UserSubscription{ + ID: 13, + UserID: 1003, + GroupID: 1, + StartsAt: start, + ExpiresAt: start.AddDate(0, 0, 30), + Status: SubscriptionStatusSuspended, + Notes: "assignment", + }) + + svc := NewSubscriptionService(groupRepo, subRepo, nil, nil, nil) + sub, err := svc.AssignSubscription(context.Background(), &AssignSubscriptionInput{ + UserID: 1003, + GroupID: 1, + ValidityDays: 30, + Notes: "assignment", + }) + + require.NoError(t, err) + require.Equal(t, int64(13), sub.ID) + require.Equal(t, SubscriptionStatusSuspended, sub.Status) + require.Equal(t, start, sub.StartsAt) + require.Equal(t, start.AddDate(0, 0, 30), sub.ExpiresAt) + require.Equal(t, "assignment", sub.Notes) + require.Equal(t, 0, subRepo.createCalls) +} + +func TestAssignSubscriptionDoesNotReactivatePastExpirySuspendedSubscription(t *testing.T) { + start := time.Now().AddDate(0, 0, -31) + expiresAt := start.AddDate(0, 0, 30) + windowStart := startOfDay(start) + groupRepo := &subscriptionGroupRepoStub{ + group: &Group{ID: 1, SubscriptionType: SubscriptionTypeSubscription}, + } + subRepo := newSubscriptionUserSubRepoStub() + subRepo.seed(&UserSubscription{ + ID: 15, + UserID: 1005, + GroupID: 1, + StartsAt: start, + ExpiresAt: expiresAt, + Status: SubscriptionStatusSuspended, + DailyWindowStart: &windowStart, + WeeklyWindowStart: &windowStart, + MonthlyWindowStart: &windowStart, + DailyUsageUSD: 1, + WeeklyUsageUSD: 2, + MonthlyUsageUSD: 3, + Notes: "suspended assignment", + }) + + svc := NewSubscriptionService(groupRepo, subRepo, nil, nil, nil) + sub, err := svc.AssignSubscription(context.Background(), &AssignSubscriptionInput{ + UserID: 1005, + GroupID: 1, + ValidityDays: 30, + Notes: "suspended assignment", + }) + + require.NoError(t, err) + require.Equal(t, int64(15), sub.ID) + require.Equal(t, SubscriptionStatusSuspended, sub.Status) + require.Equal(t, start, sub.StartsAt) + require.Equal(t, expiresAt, sub.ExpiresAt) + require.Equal(t, "suspended assignment", sub.Notes) + require.Equal(t, &windowStart, sub.DailyWindowStart) + require.Equal(t, &windowStart, sub.WeeklyWindowStart) + require.Equal(t, &windowStart, sub.MonthlyWindowStart) + require.Equal(t, float64(1), sub.DailyUsageUSD) + require.Equal(t, float64(2), sub.WeeklyUsageUSD) + require.Equal(t, float64(3), sub.MonthlyUsageUSD) + require.Equal(t, 0, subRepo.createCalls) +} + +func TestAssignSubscriptionRenewsExpiredSemanticMatch(t *testing.T) { + groupRepo := &subscriptionGroupRepoStub{ + group: &Group{ID: 1, SubscriptionType: SubscriptionTypeSubscription}, + } + subRepo := newSubscriptionUserSubRepoStub() + oldStart := time.Now().Add(-time.Hour) + oldWindowStart := startOfDay(oldStart) + subRepo.seed(&UserSubscription{ + ID: 12, + UserID: 1002, + GroupID: 1, + StartsAt: oldStart, + ExpiresAt: oldStart.AddDate(0, 0, 30), + Status: SubscriptionStatusExpired, + DailyWindowStart: &oldWindowStart, + WeeklyWindowStart: &oldWindowStart, + MonthlyWindowStart: &oldWindowStart, + DailyUsageUSD: 1, + WeeklyUsageUSD: 2, + MonthlyUsageUSD: 3, + Notes: " assignment ", + }) + + svc := NewSubscriptionService(groupRepo, subRepo, nil, nil, nil) + before := time.Now() + sub, err := svc.AssignSubscription(context.Background(), &AssignSubscriptionInput{ + UserID: 1002, + GroupID: 1, + ValidityDays: 30, + Notes: "assignment", + }) + after := time.Now() + + require.NoError(t, err) + require.Equal(t, int64(12), sub.ID) + require.Equal(t, 0, subRepo.createCalls) + require.Equal(t, SubscriptionStatusActive, sub.Status) + require.False(t, sub.StartsAt.Before(before)) + require.False(t, sub.StartsAt.After(after)) + require.Equal(t, sub.StartsAt.AddDate(0, 0, 30), sub.ExpiresAt) + require.Equal(t, startOfDay(sub.StartsAt), *sub.DailyWindowStart) + require.Equal(t, startOfDay(sub.StartsAt), *sub.WeeklyWindowStart) + require.Equal(t, startOfDay(sub.StartsAt), *sub.MonthlyWindowStart) + require.Zero(t, sub.DailyUsageUSD) + require.Zero(t, sub.WeeklyUsageUSD) + require.Zero(t, sub.MonthlyUsageUSD) + require.Equal(t, " assignment ", sub.Notes) +} + +func TestAssignSubscriptionRenewsExpiredAndAppendsDifferentNotes(t *testing.T) { + groupRepo := &subscriptionGroupRepoStub{ + group: &Group{ID: 1, SubscriptionType: SubscriptionTypeSubscription}, + } + subRepo := newSubscriptionUserSubRepoStub() + oldStart := time.Date(2025, 1, 2, 3, 4, 5, 0, time.UTC) + subRepo.seed(&UserSubscription{ + ID: 14, + UserID: 1004, + GroupID: 1, + StartsAt: oldStart, + ExpiresAt: oldStart.AddDate(0, 0, 30), + Status: SubscriptionStatusExpired, + Notes: "old assignment", + }) + + svc := NewSubscriptionService(groupRepo, subRepo, nil, nil, nil) + sub, err := svc.AssignSubscription(context.Background(), &AssignSubscriptionInput{ + UserID: 1004, + GroupID: 1, + ValidityDays: 30, + Notes: "new assignment", + }) + + require.NoError(t, err) + require.Equal(t, "old assignment\nnew assignment", sub.Notes) } func TestAssignSubscriptionConflictWhenSemanticsMismatch(t *testing.T) { - start := time.Date(2026, 2, 20, 10, 0, 0, 0, time.UTC) + start := time.Now().Add(-time.Hour) groupRepo := &subscriptionGroupRepoStub{ group: &Group{ID: 1, SubscriptionType: SubscriptionTypeSubscription}, } @@ -306,6 +466,7 @@ func TestAssignSubscriptionConflictWhenSemanticsMismatch(t *testing.T) { GroupID: 1, StartsAt: start, ExpiresAt: start.AddDate(0, 0, 30), + Status: SubscriptionStatusActive, Notes: "old-note", }) @@ -322,7 +483,7 @@ func TestAssignSubscriptionConflictWhenSemanticsMismatch(t *testing.T) { } func TestBulkAssignSubscriptionCreatedReusedAndConflict(t *testing.T) { - start := time.Date(2026, 2, 20, 10, 0, 0, 0, time.UTC) + start := time.Now().Add(-time.Hour) groupRepo := &subscriptionGroupRepoStub{ group: &Group{ID: 1, SubscriptionType: SubscriptionTypeSubscription}, } @@ -334,6 +495,7 @@ func TestBulkAssignSubscriptionCreatedReusedAndConflict(t *testing.T) { GroupID: 1, StartsAt: start, ExpiresAt: start.AddDate(0, 0, 30), + Status: SubscriptionStatusActive, Notes: "same-note", }) // user 3: 语义冲突(有效期不一致),应 failed @@ -343,6 +505,7 @@ func TestBulkAssignSubscriptionCreatedReusedAndConflict(t *testing.T) { GroupID: 1, StartsAt: start, ExpiresAt: start.AddDate(0, 0, 60), + Status: SubscriptionStatusActive, Notes: "same-note", }) @@ -365,6 +528,52 @@ func TestBulkAssignSubscriptionCreatedReusedAndConflict(t *testing.T) { require.Equal(t, 1, subRepo.createCalls) } +func TestBulkAssignSubscriptionRenewsExpiredSemanticMatch(t *testing.T) { + groupRepo := &subscriptionGroupRepoStub{ + group: &Group{ID: 1, SubscriptionType: SubscriptionTypeSubscription}, + } + subRepo := newSubscriptionUserSubRepoStub() + oldStart := time.Date(2025, 1, 2, 3, 4, 5, 0, time.UTC) + subRepo.seed(&UserSubscription{ + ID: 24, + UserID: 4, + GroupID: 1, + StartsAt: oldStart, + ExpiresAt: oldStart.AddDate(0, 0, 7), + Status: SubscriptionStatusExpired, + DailyUsageUSD: 1, + WeeklyUsageUSD: 2, + MonthlyUsageUSD: 3, + Notes: "bulk", + }) + + svc := NewSubscriptionService(groupRepo, subRepo, nil, nil, nil) + before := time.Now() + result, err := svc.BulkAssignSubscription(context.Background(), &BulkAssignSubscriptionInput{ + UserIDs: []int64{4}, + GroupID: 1, + ValidityDays: 7, + Notes: "bulk", + }) + after := time.Now() + + require.NoError(t, err) + require.Equal(t, 1, result.SuccessCount) + require.Equal(t, 0, result.CreatedCount) + require.Equal(t, 1, result.ReusedCount) + require.Equal(t, "reused", result.Statuses[4]) + require.Len(t, result.Subscriptions, 1) + renewed := result.Subscriptions[0] + require.Equal(t, SubscriptionStatusActive, renewed.Status) + require.False(t, renewed.StartsAt.Before(before)) + require.False(t, renewed.StartsAt.After(after)) + require.Equal(t, renewed.StartsAt.AddDate(0, 0, 7), renewed.ExpiresAt) + require.Zero(t, renewed.DailyUsageUSD) + require.Zero(t, renewed.WeeklyUsageUSD) + require.Zero(t, renewed.MonthlyUsageUSD) + require.Equal(t, "bulk", renewed.Notes) +} + func TestAssignSubscriptionKeepsWorkingWhenIdempotencyStoreUnavailable(t *testing.T) { groupRepo := &subscriptionGroupRepoStub{ group: &Group{ID: 1, SubscriptionType: SubscriptionTypeSubscription}, diff --git a/backend/internal/service/subscription_service.go b/backend/internal/service/subscription_service.go index ea1fd091d..fbe99bf90 100644 --- a/backend/internal/service/subscription_service.go +++ b/backend/internal/service/subscription_service.go @@ -508,6 +508,25 @@ func (s *SubscriptionService) assignSubscriptionWithReuse(ctx context.Context, i if getErr != nil { return nil, false, getErr } + now := time.Now() + if sub.Status == SubscriptionStatusExpired || + (sub.Status != SubscriptionStatusSuspended && !sub.ExpiresAt.After(now)) { + validityDays := normalizeAssignValidityDays(input.ValidityDays) + newExpiresAt := now.AddDate(0, 0, validityDays) + if newExpiresAt.After(MaxExpiresAt) { + newExpiresAt = MaxExpiresAt + } + renewalNotes := input.Notes + if strings.TrimSpace(sub.Notes) == strings.TrimSpace(input.Notes) { + renewalNotes = "" + } + if err := s.updateExistingSubscriptionTerm(ctx, sub, renewalNotes, now, newExpiresAt, true); err != nil { + return nil, false, err + } + s.maybeInvalidateAssignmentCaches(input.UserID, input.GroupID, false) + renewed, getErr := s.userSubRepo.GetByID(ctx, sub.ID) + return renewed, true, getErr + } if conflictReason, conflict := detectAssignSemanticConflict(sub, input); conflict { return nil, false, ErrSubscriptionAssignConflict.WithMetadata(map[string]string{ "conflict_reason": conflictReason, diff --git a/backend/internal/service/upstream_billing_probe.go b/backend/internal/service/upstream_billing_probe.go index 70dbf8f8f..d5c28e68b 100644 --- a/backend/internal/service/upstream_billing_probe.go +++ b/backend/internal/service/upstream_billing_probe.go @@ -38,7 +38,7 @@ const ( upstreamBillingProbeMaxBodyBytes = 64 * 1024 upstreamBillingProbeMaxPerCycle = 20 upstreamBillingProbeConcurrency = 4 - upstreamBillingProbeMaxBackoff = 24 * time.Hour + upstreamBillingProbeMaxDelay = 24 * time.Hour upstreamBillingProbeLeaderLockKey = "upstream:billing:probe:leader" upstreamBillingProbeLeaderLockTTL = 2 * time.Minute ) @@ -620,7 +620,7 @@ func (s *UpstreamBillingProbeService) probeLoadedAccount(ctx context.Context, ac ReceivedAt: probeTimePtr(now), FreshUntil: probeTimePtr(now.Add(2 * time.Duration(intervalMinutes) * time.Minute)), LastAttemptAt: now, - NextProbeAt: now.Add(nextProbeDelay(intervalMinutes, 0, 0)), + NextProbeAt: now.Add(nextProbeDelay(intervalMinutes, 0)), HTTPStatus: resp.StatusCode, } if err := s.updateSnapshot(ctx, account, snapshot); err != nil { @@ -650,7 +650,7 @@ func (s *UpstreamBillingProbeService) persistProbeFailure( snapshot := &UpstreamBillingProbeSnapshot{ Status: status, LastAttemptAt: now, - NextProbeAt: now.Add(nextProbeDelay(intervalMinutes, failureCount, retryAfterDuration)), + NextProbeAt: now.Add(nextProbeDelay(intervalMinutes, retryAfterDuration)), FailureCount: failureCount, HTTPStatus: statusCode, LastError: reason, @@ -859,20 +859,13 @@ func (s *UpstreamBillingProbeService) currentTime() time.Time { return time.Now() } -func nextProbeDelay(intervalMinutes, failureCount int, retryAfterDuration time.Duration) time.Duration { +func nextProbeDelay(intervalMinutes int, retryAfterDuration time.Duration) time.Duration { interval := time.Duration(intervalMinutes) * time.Minute if interval < upstreamBillingProbeMinIntervalMinutes*time.Minute { interval = upstreamBillingProbeMinIntervalMinutes * time.Minute } - if failureCount > 0 { - shift := failureCount - if shift > 5 { - shift = 5 - } - interval *= time.Duration(1 << shift) - } - if interval > upstreamBillingProbeMaxBackoff { - interval = upstreamBillingProbeMaxBackoff + if interval > upstreamBillingProbeMaxDelay { + interval = upstreamBillingProbeMaxDelay } jitterRange := interval / 5 if jitterRange > 5*time.Minute { @@ -883,11 +876,11 @@ func nextProbeDelay(intervalMinutes, failureCount int, retryAfterDuration time.D } if retryAfterDuration > interval { // Retry-After is an explicit upstream instruction; do not shorten it - // with the local exponential-backoff ceiling. + // with the local maximum delay. return retryAfterDuration } - if interval > upstreamBillingProbeMaxBackoff { - return upstreamBillingProbeMaxBackoff + if interval > upstreamBillingProbeMaxDelay { + return upstreamBillingProbeMaxDelay } return interval } diff --git a/backend/internal/service/upstream_billing_probe_test.go b/backend/internal/service/upstream_billing_probe_test.go index 0ffcc6006..eecd866a9 100644 --- a/backend/internal/service/upstream_billing_probe_test.go +++ b/backend/internal/service/upstream_billing_probe_test.go @@ -495,7 +495,7 @@ func TestUpstreamBillingProbeFailurePreservesLastSuccessAndRetryAfter(t *testing } func TestUpstreamBillingProbeRetryAfterIsNotShortened(t *testing.T) { - delay := nextProbeDelay(30, 1, 48*time.Hour) + delay := nextProbeDelay(30, 48*time.Hour) require.Equal(t, 48*time.Hour, delay) } @@ -510,11 +510,21 @@ func TestUpstreamBillingProbeEmptyResponseIsPersistedAsFailure(t *testing.T) { } repo := &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{account.ID: account}} svc := newUpstreamBillingProbeTestService(repo, &httpUpstreamRecorder{}, &upstreamBillingProbeSettingRepo{}) + fixedNow := time.Date(2026, time.July, 13, 2, 0, 0, 0, time.UTC) + svc.now = func() time.Time { return fixedNow } snapshot, err := svc.ProbeAccount(context.Background(), account.ID) require.NoError(t, err) require.Equal(t, UpstreamBillingProbeStatusFailed, snapshot.Status) require.Equal(t, "empty_response", snapshot.LastError) + require.False(t, snapshot.NextProbeAt.Before(fixedNow.Add(24*time.Minute))) + require.False(t, snapshot.NextProbeAt.After(fixedNow.Add(36*time.Minute))) + + snapshot, err = svc.ProbeAccount(context.Background(), account.ID) + require.NoError(t, err) + require.Equal(t, 2, snapshot.FailureCount) + require.False(t, snapshot.NextProbeAt.Before(fixedNow.Add(24*time.Minute))) + require.False(t, snapshot.NextProbeAt.After(fixedNow.Add(36*time.Minute))) } func TestUpstreamBillingProbeUnsupportedAndAccountToggle(t *testing.T) { @@ -533,6 +543,8 @@ func TestUpstreamBillingProbeUnsupportedAndAccountToggle(t *testing.T) { Body: io.NopCloser(strings.NewReader("not found")), }} svc := newUpstreamBillingProbeTestService(repo, upstream, &upstreamBillingProbeSettingRepo{}) + fixedNow := time.Date(2026, time.July, 13, 2, 0, 0, 0, time.UTC) + svc.now = func() time.Time { return fixedNow } require.NoError(t, svc.SetAccountEnabled(context.Background(), account.ID, true)) require.Equal(t, true, account.Extra[UpstreamBillingProbeEnabledExtraKey]) @@ -540,6 +552,14 @@ func TestUpstreamBillingProbeUnsupportedAndAccountToggle(t *testing.T) { require.NoError(t, err) require.Equal(t, UpstreamBillingProbeStatusUnsupported, snapshot.Status) require.Equal(t, "unsupported", snapshot.LastError) + require.False(t, snapshot.NextProbeAt.Before(fixedNow.Add(24*time.Minute))) + require.False(t, snapshot.NextProbeAt.After(fixedNow.Add(36*time.Minute))) + + snapshot, err = svc.ProbeAccount(context.Background(), account.ID) + require.NoError(t, err) + require.Equal(t, 2, snapshot.FailureCount) + require.False(t, snapshot.NextProbeAt.Before(fixedNow.Add(24*time.Minute))) + require.False(t, snapshot.NextProbeAt.After(fixedNow.Add(36*time.Minute))) invalid := &Account{ID: 20, Platform: PlatformOpenAI, Type: AccountTypeOAuth} repo.accounts[invalid.ID] = invalid diff --git a/backend/internal/service/user_subscription_daily_quota_test.go b/backend/internal/service/user_subscription_daily_quota_test.go index bf58de7f4..bc0481216 100644 --- a/backend/internal/service/user_subscription_daily_quota_test.go +++ b/backend/internal/service/user_subscription_daily_quota_test.go @@ -65,6 +65,35 @@ func TestAssignOrExtendSubscription_ExpiredDailyCardStartsNewOneTimeQuota(t *tes require.Equal(t, "old\nnew", renewed.Notes) } +func TestAssignOrExtendSubscription_ExpiredSubscriptionAppendsMatchingNotes(t *testing.T) { + groupRepo := &subscriptionGroupRepoStub{ + group: &Group{ID: 1, SubscriptionType: SubscriptionTypeSubscription}, + } + subRepo := newSubscriptionUserSubRepoStub() + oldStart := time.Now().AddDate(0, 0, -3) + subRepo.seed(&UserSubscription{ + ID: 101, + UserID: 201, + GroupID: 1, + StartsAt: oldStart, + ExpiresAt: oldStart.AddDate(0, 0, 1), + Status: SubscriptionStatusExpired, + Notes: "same", + }) + svc := NewSubscriptionService(groupRepo, subRepo, nil, nil, nil) + + renewed, reused, err := svc.AssignOrExtendSubscription(context.Background(), &AssignSubscriptionInput{ + UserID: 201, + GroupID: 1, + ValidityDays: 1, + Notes: "same", + }) + + require.NoError(t, err) + require.True(t, reused) + require.Equal(t, "same\nsame", renewed.Notes) +} + func TestUserSubscriptionNeedsDailyReset_DailyCardKeepsOneTimeQuota(t *testing.T) { start := time.Date(2026, 5, 18, 12, 0, 0, 0, time.UTC) dailyWindowStart := time.Date(2026, 5, 18, 0, 0, 0, 0, time.UTC) diff --git a/backend/internal/service/wire.go b/backend/internal/service/wire.go index 59037d188..f3ee81539 100644 --- a/backend/internal/service/wire.go +++ b/backend/internal/service/wire.go @@ -567,6 +567,8 @@ func ProvideOpsService( antigravityGatewayService *AntigravityGatewayService, systemLogSink *OpsSystemLogSink, settingService *SettingService, + authCacheInvalidationWorker *AuthCacheInvalidationWorker, + apiKeyService *APIKeyService, ) *OpsService { svc := NewOpsService( opsRepo, @@ -587,9 +589,25 @@ func ProvideOpsService( // a populated cache rather than zero defaults. Best-effort, sync-bounded. settingService.WarmOpenAIQuotaAutoPauseSettings(context.Background()) } + svc.authCacheInvalidationWorker = authCacheInvalidationWorker + svc.apiKeyService = apiKeyService + svc.StartRuntimeSettingsRefresh(context.Background()) return svc } +// ProvideOpsIngressRejectAggregator starts the bounded security aggregation +// runtime and attaches it to OpsService, which is the middleware recorder. +func ProvideOpsIngressRejectAggregator(opsRepo OpsRepository, opsService *OpsService) *OpsIngressRejectAggregator { + repo, ok := opsRepo.(OpsIngressRejectRepository) + if !ok { + return nil + } + aggregator := NewOpsIngressRejectAggregator(repo) + aggregator.Start() + opsService.SetIngressRejectAggregator(aggregator) + return aggregator +} + // ProvideSettingService wires SettingService with group reader and proxy repo. func ProvideSettingService(settingRepo SettingRepository, groupRepo GroupRepository, proxyRepo ProxyRepository, cfg *config.Config) *SettingService { svc := NewSettingService(settingRepo, cfg) @@ -647,6 +665,7 @@ var ProviderSet = wire.NewSet( NewUserService, ProvideAPIKeyService, ProvideAPIKeyAuthCacheInvalidator, + ProvideAuthCacheInvalidationWorker, NewGroupService, NewAccountService, NewProxyService, @@ -696,6 +715,7 @@ var ProviderSet = wire.NewSet( ProvideBackupService, ProvideOpsSystemLogSink, ProvideOpsService, + ProvideOpsIngressRejectAggregator, ProvideAuditLogService, ProvideOpsMetricsCollector, ProvideOpsAggregationService, diff --git a/backend/internal/web/embed_on.go b/backend/internal/web/embed_on.go index 4bd984301..516ac88dc 100644 --- a/backend/internal/web/embed_on.go +++ b/backend/internal/web/embed_on.go @@ -11,6 +11,7 @@ import ( "io" "io/fs" "net/http" + "net/url" "os" "path/filepath" "strings" @@ -209,12 +210,64 @@ func (s *FrontendServer) injectSettings(settingsJSON []byte) []byte { headClose := []byte("") result := bytes.Replace(s.baseHTML, headClose, append(script, headClose...), 1) - // Replace
+ {{ t('admin.accounts.upstreamBilling.autoProbeHint') }} +
+{{ apiKeyHint }}
+ {{ t('admin.accounts.upstreamBilling.autoProbeHint') }} +
+{{ t('admin.accounts.upstreamBilling.groupRate', { value: data.group_rate_multiplier }) }}
-- {{ t('admin.accounts.upstreamBilling.userRate', { value: data.user_rate_multiplier }) }} +
{{ t('admin.accounts.upstreamBilling.groupRate', { value: data.group_rate_multiplier }) }}
++ {{ t('admin.accounts.upstreamBilling.userRate', { value: data.user_rate_multiplier }) }} +
++ {{ + data.peak_rate_enabled + ? t('admin.accounts.upstreamBilling.peakRate', { + start: data.peak_start, + end: data.peak_end, + value: data.peak_rate_multiplier, + timezone: data.timezone + }) + : t('admin.accounts.upstreamBilling.noPeakRate') + }} +
+{{ t('admin.accounts.upstreamBilling.effectiveRate', { value: currentEffectiveRate ?? '-' }) }}
+{{ t('admin.accounts.upstreamBilling.updatedAt', { value: formatDate(snapshot?.received_at) }) }}
+ + ++ {{ t('admin.accounts.upstreamBilling.lastDetectedRate', { value: lastDetectedRate }) }} +
++ {{ t('admin.accounts.upstreamBilling.lastDetectedAt', { value: formatDate(snapshot?.received_at) }) }} +
++ {{ t('admin.accounts.upstreamBilling.elapsedSince', { value: elapsedSinceLastSuccess }) }} +
+ +{{ statusLabel || '-' }}
++ {{ t('admin.accounts.upstreamBilling.nextProbeAt', { value: formatDate(nextProbeAt) }) }}
-- {{ - data.peak_rate_enabled - ? t('admin.accounts.upstreamBilling.peakRate', { - start: data.peak_start, - end: data.peak_end, - value: data.peak_rate_multiplier, - timezone: data.timezone - }) - : t('admin.accounts.upstreamBilling.noPeakRate') - }} +
+ {{ t('admin.accounts.upstreamBilling.accountProbeState') }} + + {{ probeEnabled ? t('admin.accounts.upstreamBilling.enabled') : t('admin.accounts.upstreamBilling.disabled') }} + +
++ {{ t('admin.accounts.upstreamBilling.globalProbeState') }} + {{ t('admin.accounts.upstreamBilling.disabled') }}
-{{ t('admin.accounts.upstreamBilling.effectiveRate', { value: currentEffectiveRate ?? '-' }) }}
-{{ t('admin.accounts.upstreamBilling.updatedAt', { value: formatDate(snapshot?.received_at) }) }}
+ {{ t("admin.settings.upstreamBillingProbe.description") }} +
++ {{ t("admin.settings.upstreamBillingProbe.enabledHint") }} +
++ {{ t("admin.settings.upstreamBillingProbe.intervalHint") }} +
+- {{ t('admin.ops.settings.ignoreInvalidApiKeyErrorsHint') }} -
-