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 with custom site name so the browser tab shows it immediately + // Apply custom branding before the browser paints the static defaults. result = injectSiteTitle(result, settingsJSON) + result = injectSiteFavicon(result, settingsJSON) return result } +// injectSiteFavicon replaces the static favicon with a configured, browser-safe image URL. +func injectSiteFavicon(html, settingsJSON []byte) []byte { + var cfg struct { + SiteLogo string `json:"site_logo"` + } + if err := json.Unmarshal(settingsJSON, &cfg); err != nil { + return html + } + + logoURL := safeImageURL(cfg.SiteLogo) + if logoURL == "" { + return html + } + + linkStart := bytes.Index(html, []byte(`<link rel="icon"`)) + if linkStart == -1 { + return html + } + linkEndOffset := bytes.IndexByte(html[linkStart:], '>') + if linkEndOffset == -1 { + return html + } + linkEnd := linkStart + linkEndOffset + 1 + replacement := []byte(`<link rel="icon" href="` + htmlpkg.EscapeString(logoURL) + `" />`) + + var buf bytes.Buffer + buf.Write(html[:linkStart]) + buf.Write(replacement) + buf.Write(html[linkEnd:]) + return buf.Bytes() +} + +func safeImageURL(value string) string { + trimmed := strings.TrimSpace(value) + if trimmed == "" { + return "" + } + if strings.HasPrefix(trimmed, "/") && !strings.HasPrefix(trimmed, "//") { + return trimmed + } + if strings.HasPrefix(strings.ToLower(trimmed), "data:image/") { + return trimmed + } + + parsed, err := url.Parse(trimmed) + if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") || parsed.Host == "" { + return "" + } + return trimmed +} + // injectSiteTitle replaces the static <title> in HTML with the configured site name. // This ensures the browser tab shows the correct title before JS executes. func injectSiteTitle(html, settingsJSON []byte) []byte { diff --git a/backend/internal/web/embed_test.go b/backend/internal/web/embed_test.go index 769b856ac..3a9bc1151 100644 --- a/backend/internal/web/embed_test.go +++ b/backend/internal/web/embed_test.go @@ -114,6 +114,41 @@ func TestInjectSiteTitle(t *testing.T) { }) } +func TestInjectSiteFavicon(t *testing.T) { + t.Run("replaces_favicon_with_site_logo", func(t *testing.T) { + html := []byte(`<html><head><link rel="icon" type="image/png" href="/logo.png" /></head></html>`) + settingsJSON := []byte(`{"site_logo":"https://example.com/custom-logo.png"}`) + + result := injectSiteFavicon(html, settingsJSON) + + assert.Contains(t, string(result), `<link rel="icon" href="https://example.com/custom-logo.png" />`) + assert.NotContains(t, string(result), `/logo.png`) + }) + + t.Run("supports_relative_and_data_image_urls", func(t *testing.T) { + html := []byte(`<link rel="icon" href="/logo.png" />`) + + assert.Contains(t, string(injectSiteFavicon(html, []byte(`{"site_logo":"/uploads/logo.svg"}`))), `/uploads/logo.svg`) + assert.Contains(t, string(injectSiteFavicon(html, []byte(`{"site_logo":"data:image/png;base64,abc"}`))), `data:image/png;base64,abc`) + }) + + t.Run("rejects_unsafe_logo_urls", func(t *testing.T) { + html := []byte(`<link rel="icon" href="/logo.png" />`) + + result := injectSiteFavicon(html, []byte(`{"site_logo":"javascript:alert(1)"}`)) + + assert.Equal(t, string(html), string(result)) + }) + + t.Run("escapes_logo_url_for_html", func(t *testing.T) { + html := []byte(`<link rel="icon" href="/logo.png" />`) + + result := injectSiteFavicon(html, []byte(`{"site_logo":"https://example.com/logo.png?a=1&b=2"}`)) + + assert.Contains(t, string(result), `a=1&b=2`) + }) +} + func TestReplaceNoncePlaceholder(t *testing.T) { t.Run("replaces_single_placeholder", func(t *testing.T) { html := []byte(`<script nonce="__CSP_NONCE_VALUE__">console.log('test');</script>`) diff --git a/backend/migrations/183_ops_ingress_reject_aggregates.sql b/backend/migrations/183_ops_ingress_reject_aggregates.sql new file mode 100644 index 000000000..71da4fb2e --- /dev/null +++ b/backend/migrations/183_ops_ingress_reject_aggregates.sql @@ -0,0 +1,27 @@ +SET LOCAL lock_timeout = '5s'; +SET LOCAL statement_timeout = '10min'; + +CREATE TABLE IF NOT EXISTS ops_ingress_reject_aggregates ( + id BIGSERIAL PRIMARY KEY, + bucket_start TIMESTAMPTZ NOT NULL, + reject_reason VARCHAR(64) NOT NULL, + route_family VARCHAR(64) NOT NULL, + protocol VARCHAR(32) NOT NULL, + client_ip INET NOT NULL, + user_id BIGINT NOT NULL DEFAULT 0, + api_key_id BIGINT NOT NULL DEFAULT 0, + request_count BIGINT NOT NULL DEFAULT 0, + first_seen TIMESTAMPTZ NOT NULL, + last_seen TIMESTAMPTZ NOT NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + CONSTRAINT ops_ingress_reject_aggregates_dimensions_unique UNIQUE + (bucket_start, reject_reason, route_family, protocol, client_ip, user_id, api_key_id) +); + +CREATE INDEX IF NOT EXISTS idx_ops_ingress_reject_aggregates_bucket + ON ops_ingress_reject_aggregates (bucket_start DESC); +CREATE INDEX IF NOT EXISTS idx_ops_ingress_reject_aggregates_reason_bucket + ON ops_ingress_reject_aggregates (reject_reason, bucket_start DESC); +CREATE INDEX IF NOT EXISTS idx_ops_ingress_reject_aggregates_ip_bucket + ON ops_ingress_reject_aggregates (client_ip, bucket_start DESC); diff --git a/backend/migrations/184_auth_cache_invalidation_outbox.sql b/backend/migrations/184_auth_cache_invalidation_outbox.sql new file mode 100644 index 000000000..a882f1079 --- /dev/null +++ b/backend/migrations/184_auth_cache_invalidation_outbox.sql @@ -0,0 +1,198 @@ +-- Durable, transactionally-enqueued API-key auth cache invalidation. +-- cache_key is always SHA-256 hex; plaintext credentials never leave api_keys. + +CREATE TABLE IF NOT EXISTS auth_cache_invalidation_outbox ( + id BIGSERIAL PRIMARY KEY, + cache_key CHAR(64) NOT NULL CHECK (cache_key ~ '^[0-9a-f]{64}$'), + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + available_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + delivery_stage SMALLINT NOT NULL DEFAULT 0 CHECK (delivery_stage IN (0, 1)), + attempts INTEGER NOT NULL DEFAULT 0 CHECK (attempts >= 0), + last_error TEXT, + claimed_at TIMESTAMPTZ, + claimed_by TEXT +); + +CREATE INDEX IF NOT EXISTS idx_auth_cache_invalidation_outbox_available + ON auth_cache_invalidation_outbox (available_at, id) + WHERE claimed_at IS NULL; +CREATE INDEX IF NOT EXISTS idx_auth_cache_invalidation_outbox_lease + ON auth_cache_invalidation_outbox (claimed_at) + WHERE claimed_at IS NOT NULL; +CREATE INDEX IF NOT EXISTS idx_auth_cache_invalidation_outbox_cache_key + ON auth_cache_invalidation_outbox (cache_key); +CREATE INDEX IF NOT EXISTS idx_auth_cache_invalidation_outbox_created_at + ON auth_cache_invalidation_outbox (created_at); + +CREATE OR REPLACE FUNCTION enqueue_auth_cache_invalidation(raw_key TEXT) +RETURNS VOID +LANGUAGE plpgsql +AS $$ +BEGIN + IF raw_key IS NULL OR raw_key = '' THEN + RETURN; + END IF; + INSERT INTO auth_cache_invalidation_outbox (cache_key) + VALUES (encode(sha256(convert_to(raw_key, 'UTF8')), 'hex')); +END; +$$; + +CREATE OR REPLACE FUNCTION enqueue_api_key_auth_cache_invalidation() +RETURNS TRIGGER +LANGUAGE plpgsql +AS $$ +BEGIN + IF TG_OP = 'DELETE' THEN + PERFORM enqueue_auth_cache_invalidation(OLD.key); + RETURN OLD; + END IF; + + IF OLD.key IS DISTINCT FROM NEW.key + OR OLD.status IS DISTINCT FROM NEW.status + OR OLD.deleted_at IS DISTINCT FROM NEW.deleted_at + OR OLD.user_id IS DISTINCT FROM NEW.user_id + OR OLD.group_id IS DISTINCT FROM NEW.group_id + OR OLD.ip_whitelist IS DISTINCT FROM NEW.ip_whitelist + OR OLD.ip_blacklist IS DISTINCT FROM NEW.ip_blacklist + OR OLD.expires_at IS DISTINCT FROM NEW.expires_at THEN + PERFORM enqueue_auth_cache_invalidation(OLD.key); + IF NEW.deleted_at IS NULL AND NEW.key IS DISTINCT FROM OLD.key THEN + PERFORM enqueue_auth_cache_invalidation(NEW.key); + END IF; + END IF; + RETURN NEW; +END; +$$; + +DROP TRIGGER IF EXISTS trg_api_keys_auth_cache_invalidation ON api_keys; +CREATE TRIGGER trg_api_keys_auth_cache_invalidation +AFTER UPDATE OR DELETE ON api_keys +FOR EACH ROW EXECUTE FUNCTION enqueue_api_key_auth_cache_invalidation(); + +CREATE OR REPLACE FUNCTION enqueue_user_auth_cache_invalidation() +RETURNS TRIGGER +LANGUAGE plpgsql +AS $$ +DECLARE + target_user_id BIGINT; +BEGIN + target_user_id := OLD.id; + IF TG_OP = 'UPDATE' + AND OLD.status IS NOT DISTINCT FROM NEW.status + AND OLD.role IS NOT DISTINCT FROM NEW.role + AND OLD.deleted_at IS NOT DISTINCT FROM NEW.deleted_at THEN + RETURN NEW; + END IF; + + INSERT INTO auth_cache_invalidation_outbox (cache_key) + SELECT encode(sha256(convert_to(k.key, 'UTF8')), 'hex') + FROM api_keys AS k + WHERE k.user_id = target_user_id + AND k.deleted_at IS NULL + AND k.key <> ''; + IF TG_OP = 'DELETE' THEN + RETURN OLD; + END IF; + RETURN NEW; +END; +$$; + +DROP TRIGGER IF EXISTS trg_users_auth_cache_invalidation ON users; +CREATE TRIGGER trg_users_auth_cache_invalidation +AFTER UPDATE OR DELETE ON users +FOR EACH ROW EXECUTE FUNCTION enqueue_user_auth_cache_invalidation(); + +CREATE OR REPLACE FUNCTION enqueue_group_auth_cache_invalidation() +RETURNS TRIGGER +LANGUAGE plpgsql +AS $$ +DECLARE + target_group_id BIGINT; +BEGIN + target_group_id := OLD.id; + IF TG_OP = 'UPDATE' + AND OLD.status IS NOT DISTINCT FROM NEW.status + AND OLD.is_exclusive IS NOT DISTINCT FROM NEW.is_exclusive + AND OLD.deleted_at IS NOT DISTINCT FROM NEW.deleted_at THEN + RETURN NEW; + END IF; + + INSERT INTO auth_cache_invalidation_outbox (cache_key) + SELECT encode(sha256(convert_to(k.key, 'UTF8')), 'hex') + FROM api_keys AS k + WHERE k.group_id = target_group_id + AND k.deleted_at IS NULL + AND k.key <> ''; + IF TG_OP = 'DELETE' THEN + RETURN OLD; + END IF; + RETURN NEW; +END; +$$; + +DROP TRIGGER IF EXISTS trg_groups_auth_cache_invalidation ON groups; +CREATE TRIGGER trg_groups_auth_cache_invalidation +AFTER UPDATE OR DELETE ON groups +FOR EACH ROW EXECUTE FUNCTION enqueue_group_auth_cache_invalidation(); + +CREATE OR REPLACE FUNCTION enqueue_allowed_group_auth_cache_invalidation() +RETURNS TRIGGER +LANGUAGE plpgsql +AS $$ +DECLARE + target_user_id BIGINT; + target_group_id BIGINT; +BEGIN + IF TG_OP = 'UPDATE' + AND (OLD.user_id IS DISTINCT FROM NEW.user_id + OR OLD.group_id IS DISTINCT FROM NEW.group_id) THEN + IF EXISTS ( + SELECT 1 FROM groups g + WHERE g.id = OLD.group_id AND g.is_exclusive = TRUE + ) THEN + INSERT INTO auth_cache_invalidation_outbox (cache_key) + SELECT encode(sha256(convert_to(k.key, 'UTF8')), 'hex') + FROM api_keys AS k + WHERE k.user_id = OLD.user_id + AND k.group_id = OLD.group_id + AND k.deleted_at IS NULL + AND k.key <> ''; + END IF; + target_user_id := NEW.user_id; + target_group_id := NEW.group_id; + ELSIF TG_OP = 'UPDATE' THEN + RETURN NEW; + ELSIF TG_OP = 'INSERT' THEN + target_user_id := NEW.user_id; + target_group_id := NEW.group_id; + ELSE + target_user_id := OLD.user_id; + target_group_id := OLD.group_id; + END IF; + + IF EXISTS ( + SELECT 1 FROM groups g + WHERE g.id = target_group_id AND g.is_exclusive = TRUE + ) THEN + INSERT INTO auth_cache_invalidation_outbox (cache_key) + SELECT encode(sha256(convert_to(k.key, 'UTF8')), 'hex') + FROM api_keys AS k + WHERE k.user_id = target_user_id + AND k.group_id = target_group_id + AND k.deleted_at IS NULL + AND k.key <> ''; + END IF; + IF TG_OP = 'DELETE' THEN + RETURN OLD; + END IF; + RETURN NEW; +END; +$$; + +DROP TRIGGER IF EXISTS trg_user_allowed_groups_auth_cache_invalidation ON user_allowed_groups; +CREATE TRIGGER trg_user_allowed_groups_auth_cache_invalidation +AFTER INSERT OR UPDATE OR DELETE ON user_allowed_groups +FOR EACH ROW EXECUTE FUNCTION enqueue_allowed_group_auth_cache_invalidation(); + +COMMENT ON TABLE auth_cache_invalidation_outbox IS + 'Durable cross-instance auth cache invalidations; cache_key is SHA-256 hex, never plaintext API key'; diff --git a/backend/scripts/finalize-ingress-reject-cleanup.sql b/backend/scripts/finalize-ingress-reject-cleanup.sql new file mode 100644 index 000000000..91a4e1ac8 --- /dev/null +++ b/backend/scripts/finalize-ingress-reject-cleanup.sql @@ -0,0 +1,23 @@ +-- Post-rollout finalizer for ingress-rejection log cleanup. +-- +-- DO NOT run this during a rolling deployment. Run it only after: +-- 1. every application instance is on the release that no longer reads or +-- writes deleted_api_key_audits and the deprecated ops_error_logs columns; +-- 2. cleanup-ingress-reject-logs has been dry-run and, if desired, executed; +-- 3. a database backup or recovery point has been verified. + +BEGIN; +SET LOCAL lock_timeout = '5s'; +SET LOCAL statement_timeout = '10min'; + +DROP TABLE IF EXISTS deleted_api_key_audits; + +ALTER TABLE IF EXISTS ops_error_logs + DROP COLUMN IF EXISTS attempted_key_prefix, + DROP COLUMN IF EXISTS deleted_key_owner_user_id, + DROP COLUMN IF EXISTS deleted_key_name; + +COMMIT; + +-- Run this separately during a normal maintenance window: +-- VACUUM (ANALYZE) ops_error_logs; diff --git a/deploy/Caddyfile b/deploy/Caddyfile index 7da36ac8c..079b5f07e 100644 --- a/deploy/Caddyfile +++ b/deploy/Caddyfile @@ -1,12 +1,25 @@ +{ + servers { + max_header_size 64KB + timeouts { + read_header 10s + idle 2m + } + } +} + # 修改为你的域名 api.sub2api.com { + # This baseline assumes clients connect directly to Caddy. When Caddy is + # behind a CDN, configure explicit trusted proxy CIDRs and {client_ip} as + # documented in EDGE_SECURITY.md; {remote_host} would otherwise be the CDN. # ========================================================================= # TLS 安全配置 # ========================================================================= tls { # 仅使用 TLS 1.2 和 1.3 protocols tls1.2 tls1.3 - + # 优先使用的加密套件 ciphers TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384 TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384 TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305_SHA256 TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305_SHA256 TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256 TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256 } @@ -20,22 +33,17 @@ api.sub2api.com { health_interval 30s health_timeout 10s health_status 200 - + # 负载均衡策略(单节点可忽略,多节点时有用) lb_policy round_robin lb_try_duration 5s lb_try_interval 250ms - - # 传递真实客户端信息 - # 兼容 Cloudflare 和直连:后端应优先读取 CF-Connecting-IP,其次 X-Real-IP + + # 仅从实际 TCP 对端生成转发头,避免透传客户端伪造值。 header_up X-Real-IP {remote_host} header_up X-Forwarded-For {remote_host} header_up X-Forwarded-Proto {scheme} header_up X-Forwarded-Host {host} - # 保留 Cloudflare 原始头(如果存在) - # 后端获取 IP 的优先级建议: CF-Connecting-IP → X-Real-IP → X-Forwarded-For - header_up CF-Connecting-IP {http.request.header.CF-Connecting-IP} - # 连接池优化 transport http { keepalive 120s @@ -44,7 +52,7 @@ api.sub2api.com { write_buffer 16KB compression off } - + # 故障转移 fail_duration 30s max_fails 3 @@ -72,7 +80,7 @@ api.sub2api.com { # 请求大小限制 (防止大文件攻击) # ========================================================================= request_body { - max_size 100MB + max_size 256MB } # ========================================================================= diff --git a/deploy/EDGE_SECURITY.md b/deploy/EDGE_SECURITY.md new file mode 100644 index 000000000..0054253eb --- /dev/null +++ b/deploy/EDGE_SECURITY.md @@ -0,0 +1,155 @@ +# Edge and HTTP Ingress Security + +Sub2API supports long-lived SSE and WebSocket requests. Protect the request +ingress without imposing a response `WriteTimeout`: a write deadline would +terminate healthy long generations and streams. + +## Application defaults + +- `server.max_header_bytes: 65536` limits HTTP/1 request headers to 64 KiB; + Go maps it to the corresponding HTTP/2 header-list limit. +- `server.read_header_timeout: 10` bounds slow-header attacks. It does not + limit request processing or response streaming. +- `server.max_request_body_size: 268435456` is the absolute 256 MiB safety net. +- `gateway.max_body_size: 268435456` remains available to multimodal, Gemini, + image, video, and batch-image endpoints. +- `gateway.text_max_body_size: 33554432` limits the known pure-text + `/embeddings` and `/alpha/search` endpoints to 32 MiB. +- H2C defaults to 50 concurrent streams per connection, a 2 MiB connection + upload window, and a 512 KiB stream upload window. +- Invalid credential abuse is limited in process by trusted client IP (IPv6 + `/64`): 120 failures per 60 seconds followed by a 60-second block. This is a + per-instance safety net; multi-instance enforcement still belongs at the + load balancer, CDN, or WAF. + +Do not add a single application-wide request semaphore: an SSE request may +legitimately occupy it for many minutes. Apply connection and unauthenticated +request controls at the edge; authenticated user/API-key concurrency remains +the application's responsibility. + +## Trusted client IPs + +`server.trusted_proxies` must contain only the CIDR/IP addresses that connect +directly to Sub2API, normally the local Nginx/Caddy address or the private load +balancer subnet. An empty list disables forwarded-IP trust. + +Never trust `CF-Connecting-IP`, `X-Real-IP`, or `X-Forwarded-For` merely because +the header exists. A CDN deployment must firewall the origin so only the CDN or +load balancer can reach it, and the proxy must overwrite forwarded headers. + +Example for a proxy on the same host: + +```yaml +server: + trusted_proxies: + - 127.0.0.1/32 + - ::1/128 +``` + +## Nginx baseline + +Define shared zones in the `http` block. Tune rates to measured legitimate +traffic; the values below are conservative starting points, not universal +capacity targets. + +```nginx +limit_conn_zone $binary_remote_addr zone=sub2api_conn:20m; +limit_req_zone $binary_remote_addr zone=sub2api_auth:20m rate=5r/s; +limit_req_zone $binary_remote_addr zone=sub2api_api:40m rate=30r/s; +map $http_upgrade $connection_upgrade { + default upgrade; + '' close; +} + +server { + listen 443 ssl http2; + server_name api.example.com; + + client_header_timeout 10s; + client_max_body_size 256m; + large_client_header_buffers 4 16k; + limit_conn sub2api_conn 40; + + location ~ ^/(auth|api/auth)/ { + limit_req zone=sub2api_auth burst=10 nodelay; + proxy_pass http://127.0.0.1:8080; + } + + location ~ ^/(v1/)?(embeddings|alpha/search)$ { + client_max_body_size 32m; + limit_req zone=sub2api_api burst=60 nodelay; + proxy_pass http://127.0.0.1:8080; + } + + location / { + limit_req zone=sub2api_api burst=60 nodelay; + proxy_http_version 1.1; + proxy_set_header Host $host; + proxy_set_header X-Real-IP $remote_addr; + proxy_set_header X-Forwarded-For $remote_addr; + proxy_set_header X-Forwarded-Proto $scheme; + proxy_set_header Upgrade $http_upgrade; + proxy_set_header Connection $connection_upgrade; + proxy_buffering off; + proxy_request_buffering off; + proxy_read_timeout 1800s; + proxy_send_timeout 1800s; + proxy_pass http://127.0.0.1:8080; + } +} +``` + +Do not use an incoming `$http_x_forwarded_for` value unless Nginx real-IP +processing is restricted to explicit trusted proxy CIDRs. + +## Caddy and CDN + +The bundled `deploy/Caddyfile` sets a 64 KiB header limit, a 10-second header +timeout, a 256 MiB absolute body limit, and overwrites forwarded addresses from +the TCP peer. It is therefore a direct-to-Caddy baseline. Do not use its +`{remote_host}` forwarding lines unchanged behind a CDN: all clients would be +attributed to a CDN egress address, collapsing rejection aggregation and the +invalid-auth limiter onto unrelated users. + +For a CDN deployment, first firewall the origin so only current CDN egress +CIDRs can connect. Then configure those exact ranges as Caddy trusted proxies +and derive upstream headers from Caddy's parsed `{client_ip}`. For example: + +```caddyfile +{ + servers { + trusted_proxies static 192.0.2.0/24 2001:db8:1234::/48 + trusted_proxies_strict + client_ip_headers CF-Connecting-IP X-Forwarded-For + } +} + +api.example.com { + reverse_proxy 127.0.0.1:8080 { + header_up X-Real-IP {client_ip} + header_up X-Forwarded-For {client_ip} + } +} +``` + +Replace the documentation ranges with the CDN's published, automatically +maintained egress ranges. `CF-Connecting-IP` is safe here only because direct +origin access is blocked and Caddy trusts only those TCP peers. Configure +Sub2API `server.trusted_proxies` with the Caddy address/private subnet so the +application accepts only Caddy's rewritten headers. + +Caddy core does not provide a general request-rate limiter; use a trusted +CDN/WAF, a supported rate-limit module, or host firewall controls. + +At a CDN/WAF, configure connection limits, header/body limits, bot challenges, +and per-IP/ASN rates before traffic reaches the origin. Allow origin ingress +only from CDN egress CIDRs or a private load balancer. Keep the application port +off the public Internet. + +## DDoS boundary + +Application checks reduce amplification after a connection reaches Go. They +cannot absorb volumetric attacks, TLS floods, bandwidth saturation, or a large +distributed source set. Those require upstream network capacity, CDN/WAF +filtering, provider firewall rules, and origin isolation. Avoid high-cardinality +metrics or per-request database security logs during rejection storms. diff --git a/deploy/README.md b/deploy/README.md index d4fcc133b..84050b188 100644 --- a/deploy/README.md +++ b/deploy/README.md @@ -27,6 +27,7 @@ This directory contains files for deploying Sub2API on Linux servers and Apple-s | `sub2api-datamanagementd.service` | datamanagementd systemd service unit file | | `DATAMANAGEMENTD_CN.md` | datamanagementd 部署与联动说明(中文) | | `config.example.yaml` | Example configuration file | +| `EDGE_SECURITY.md` | Reverse proxy, CDN/WAF, trusted proxy, and ingress hardening guide | --- diff --git a/deploy/config.example.yaml b/deploy/config.example.yaml index eff3f401b..24fa44515 100644 --- a/deploy/config.example.yaml +++ b/deploy/config.example.yaml @@ -27,6 +27,15 @@ server: # 用于生成邮件中的外部链接(例如:重置密码链接)的前端基础地址 # Example: "https://example.com" frontend_url: "" + # Maximum time to receive complete request headers. Does not limit response streams. + # 完整读取请求头的最大时间;不限制响应流持续时间。 + read_header_timeout: 10 + # Request header limit in bytes (64 KiB); also bounds the HTTP/2 header list. + # 请求头上限(字节,默认 64 KiB);同时约束 HTTP/2 header list。 + max_header_bytes: 65536 + # Keep-alive idle timeout in seconds. + # Keep-Alive 空闲连接超时(秒)。 + idle_timeout: 120 # Trusted proxies for X-Forwarded-For parsing (CIDR/IP). Empty disables trusted proxies. # 信任的代理地址(CIDR/IP 格式),用于解析 X-Forwarded-For 头。留空则禁用代理信任。 trusted_proxies: [] @@ -167,6 +176,9 @@ gateway: # Max request body size in bytes (default: 256MB) # 请求体最大字节数(默认 256MB) max_body_size: 268435456 + # Pure-text endpoint body limit (embeddings and alpha/search), default 32 MiB. + # 纯文本端点请求体上限(embeddings、alpha/search),默认 32 MiB。 + text_max_body_size: 33554432 # Max bytes to read for non-stream upstream responses (default: 8MB) # 非流式上游响应体读取上限(默认 8MB) upstream_response_read_max_bytes: 8388608 @@ -261,9 +273,10 @@ gateway: openai_compact_model: "gpt-5.4" # OpenAI Responses WebSocket 配置(默认开启,可按需回滚到 HTTP) openai_ws: - # 新版 WS mode 路由(默认关闭)。关闭时保持当前 legacy 实现行为。 + # 新版 WS mode 路由(默认关闭)。关闭时忽略账号级 WS mode(包括 http_bridge),保持 legacy ctx_pool 行为。 + # 环境变量:GATEWAY_OPENAI_WS_MODE_ROUTER_V2_ENABLED=true mode_router_v2_enabled: false - # ingress 默认模式:off|ctx_pool|passthrough|http_bridge(仅 mode_router_v2_enabled=true 生效) + # ingress 默认模式:off|ctx_pool|passthrough|http_bridge(仅 mode_router_v2_enabled=true 生效;关闭时账号级 WS mode 被忽略) # 兼容旧值:shared/dedicated 会按 ctx_pool 处理。 ingress_mode_default: ctx_pool # 完整读取并解压首条客户端消息的总超时(秒);大请求或慢链路可调高到 120-300。 @@ -708,6 +721,23 @@ api_key_auth_cache: # Enable singleflight for cache misses # 缓存未命中时启用 singleflight 合并回源 singleflight: true + # Maximum concurrent database lookups for authentication cache misses + # 认证缓存未命中时允许并发回源数据库的最大数量 + lookup_concurrency: 64 + # Process-local invalid-auth abuse protection. Counts only missing, malformed, + # deprecated-query, and confirmed invalid credentials; valid requests and + # Redis/DB failures do not consume the budget. + # 本机无效鉴权防护:仅统计缺失、格式错误、废弃 query 及确认无效的凭据。 + invalid_abuse: + enabled: true + # Invalid attempts per trusted client IP (IPv6 grouped by /64) per window. + # 每个可信客户端 IP(IPv6 按 /64 聚合)在窗口内允许的无效次数。 + threshold: 120 + window_seconds: 60 + block_seconds: 60 + # Maximum tracked client identities per process; memory remains bounded. + # 每进程最多跟踪的客户端身份数量,确保内存有界。 + capacity: 16384 # ============================================================================= # Dashboard Cache Configuration diff --git a/deploy/docker-compose.yml b/deploy/docker-compose.yml index 637c2762f..e60776e58 100644 --- a/deploy/docker-compose.yml +++ b/deploy/docker-compose.yml @@ -241,12 +241,16 @@ services: hard: 100000 volumes: - redis_data:/data + # The command is one quoted script for the inner `sh -c`. Compose keeps + # the newlines inside the quoted string, so every line needs a trailing + # `\` — without it, `redis-server` on the first line runs with no flags + # at all, and the --save/--appendonly/--appendfsync lines are never read. command: > sh -c ' - redis-server - --save 60 1 - --appendonly yes - --appendfsync everysec + redis-server \ + --save 60 1 \ + --appendonly yes \ + --appendfsync everysec \ ${REDIS_PASSWORD:+--requirepass "$REDIS_PASSWORD"}' environment: - TZ=${TZ:-Asia/Shanghai} diff --git a/frontend/src/App.vue b/frontend/src/App.vue index 14c364c3b..1a4f7437d 100644 --- a/frontend/src/App.vue +++ b/frontend/src/App.vue @@ -8,6 +8,7 @@ import { resolveRouteDocumentTitle } from '@/router/title' import AnnouncementPopup from '@/components/common/AnnouncementPopup.vue' import { useAppStore, useAuthStore, useSubscriptionStore, useAnnouncementStore, useAdminComplianceStore, useAdminSettingsStore } from '@/stores' import { getSetupStatus } from '@/api/setup' +import { updateFavicon } from '@/utils/branding' const router = useRouter() const route = useRoute() @@ -26,22 +27,6 @@ function updateDocumentTitle() { document.title = resolveRouteDocumentTitle(route, appStore.siteName, customMenuItems) } -/** - * Update favicon dynamically - * @param logoUrl - URL of the logo to use as favicon - */ -function updateFavicon(logoUrl: string) { - // Find existing favicon link or create new one - let link = document.querySelector<HTMLLinkElement>('link[rel="icon"]') - if (!link) { - link = document.createElement('link') - link.rel = 'icon' - document.head.appendChild(link) - } - link.type = logoUrl.endsWith('.svg') ? 'image/svg+xml' : 'image/x-icon' - link.href = logoUrl -} - // Watch for site settings changes and update favicon/title watch( () => appStore.siteLogo, diff --git a/frontend/src/api/admin/ops.ts b/frontend/src/api/admin/ops.ts index 3284ef7e1..bb8b31eb2 100644 --- a/frontend/src/api/admin/ops.ts +++ b/frontend/src/api/admin/ops.ts @@ -935,10 +935,6 @@ export interface OpsErrorLog { request_type?: number | null user_agent?: string - // 已删除 KEY 所有者(INVALID_API_KEY 归因快照):认证失败行 user_id 为空, - // 用户列以此回退显示所有者 - deleted_key_owner_user_id?: number | null - deleted_key_owner_email?: string | null } export interface OpsErrorDetail extends OpsErrorLog { @@ -958,11 +954,6 @@ export interface OpsErrorDetail extends OpsErrorLog { is_business_limited: boolean - // Deleted key owner info (INVALID_API_KEY attribution); - // owner user_id/email 已上移到 OpsErrorLog(列表用户列回退) - attempted_key_prefix?: string | null - deleted_key_name?: string | null - // Bound (non-deleted) key prefix, snapshotted at error time api_key_prefix?: string | null } diff --git a/frontend/src/components/account/BulkEditAccountModal.vue b/frontend/src/components/account/BulkEditAccountModal.vue index 5f744c0fc..a4268fb92 100644 --- a/frontend/src/components/account/BulkEditAccountModal.vue +++ b/frontend/src/components/account/BulkEditAccountModal.vue @@ -854,6 +854,45 @@ </div> </div> + <!-- Upstream billing auto probe (OpenAI API Key only) --> + <div v-if="allOpenAIAPIKey" class="border-t border-gray-200 pt-4 dark:border-dark-600"> + <div class="mb-3 flex items-center justify-between"> + <div class="flex-1 pr-4"> + <label + id="bulk-edit-upstream-billing-auto-probe-label" + class="input-label mb-0" + for="bulk-edit-upstream-billing-auto-probe-enabled" + > + {{ t('admin.accounts.upstreamBilling.autoProbe') }} + </label> + <p class="mt-1 text-xs text-gray-500 dark:text-gray-400"> + {{ t('admin.accounts.upstreamBilling.autoProbeHint') }} + </p> + </div> + <input + v-model="enableUpstreamBillingAutoProbe" + id="bulk-edit-upstream-billing-auto-probe-enabled" + type="checkbox" + aria-controls="bulk-edit-upstream-billing-auto-probe" + class="rounded border-gray-300 text-primary-600 focus:ring-primary-500" + /> + </div> + <div + id="bulk-edit-upstream-billing-auto-probe" + :class="!enableUpstreamBillingAutoProbe && 'pointer-events-none opacity-50'" + role="group" + aria-labelledby="bulk-edit-upstream-billing-auto-probe-label" + > + <Select + v-model="upstreamBillingAutoProbeMode" + :disabled="!enableUpstreamBillingAutoProbe" + data-testid="bulk-edit-upstream-billing-auto-probe-select" + :options="upstreamBillingAutoProbeOptions" + aria-labelledby="bulk-edit-upstream-billing-auto-probe-label" + /> + </div> + </div> + <!-- OpenAI API Key WS mode --> <div v-if="allOpenAIAPIKey" class="border-t border-gray-200 pt-4 dark:border-dark-600"> <div class="mb-3 flex items-center justify-between"> @@ -1361,6 +1400,7 @@ const enableGroups = ref(false) const enableOpenAIPassthrough = ref(false) const enableOpenAIWSMode = ref(false) const enableOpenAIAPIKeyWSMode = ref(false) +const enableUpstreamBillingAutoProbe = ref(false) const enableCodexCLIOnly = ref(false) const enableCodexCLIOnlyAppServer = ref(false) const enableOpenAICompactMode = ref(false) @@ -1391,6 +1431,7 @@ const groupIds = ref<number[]>([]) const openaiPassthroughEnabled = ref(false) const openaiOAuthResponsesWebSocketV2Mode = ref<OpenAIWSMode>(OPENAI_WS_MODE_OFF) const openaiAPIKeyResponsesWebSocketV2Mode = ref<OpenAIWSMode>(OPENAI_WS_MODE_OFF) +const upstreamBillingAutoProbeMode = ref<'enabled' | 'disabled'>('enabled') const codexCLIOnlyEnabled = ref(false) const codexCLIOnlyAppServerEnabled = ref(false) const openAICompactMode = ref<OpenAICompactMode>('auto') @@ -1421,6 +1462,10 @@ const statusOptions = computed(() => [ { value: 'active', label: t('common.active') }, { value: 'inactive', label: t('common.inactive') } ]) +const upstreamBillingAutoProbeOptions = computed(() => [ + { value: 'enabled', label: t('common.enabled') }, + { value: 'disabled', label: t('common.disabled') } +]) const isOpenAIModelRestrictionDisabled = computed( () => allOpenAIPassthroughCapable.value && @@ -1648,6 +1693,10 @@ const buildUpdatePayload = (): Record<string, unknown> | null => { ) } + if (enableUpstreamBillingAutoProbe.value) { + updates.upstream_billing_probe_enabled = upstreamBillingAutoProbeMode.value === 'enabled' + } + if (enableCodexCLIOnly.value) { const extra = ensureExtra() extra.codex_cli_only = codexCLIOnlyEnabled.value @@ -1770,6 +1819,7 @@ const handleSubmit = async () => { enableGroups.value || enableOpenAIWSMode.value || enableOpenAIAPIKeyWSMode.value || + enableUpstreamBillingAutoProbe.value || enableCodexCLIOnly.value || enableCodexCLIOnlyAppServer.value || enableOpenAICompactMode.value || @@ -1898,6 +1948,7 @@ watch( enableOpenAIPassthrough.value = false enableOpenAIWSMode.value = false enableOpenAIAPIKeyWSMode.value = false + enableUpstreamBillingAutoProbe.value = false enableCodexCLIOnly.value = false enableCodexCLIOnlyAppServer.value = false enableOpenAICompactMode.value = false @@ -1924,6 +1975,7 @@ watch( groupIds.value = [] openaiOAuthResponsesWebSocketV2Mode.value = OPENAI_WS_MODE_OFF openaiAPIKeyResponsesWebSocketV2Mode.value = OPENAI_WS_MODE_OFF + upstreamBillingAutoProbeMode.value = 'enabled' codexCLIOnlyEnabled.value = false codexCLIOnlyAppServerEnabled.value = false openAICompactMode.value = 'auto' diff --git a/frontend/src/components/account/CreateAccountModal.vue b/frontend/src/components/account/CreateAccountModal.vue index 7eebce530..f71bbbada 100644 --- a/frontend/src/components/account/CreateAccountModal.vue +++ b/frontend/src/components/account/CreateAccountModal.vue @@ -1143,6 +1143,23 @@ <p v-if="apiKeyHint" class="input-hint">{{ apiKeyHint }}</p> </div> + <div + v-if="form.platform === 'openai'" + class="flex items-center justify-between gap-4 border-t border-gray-200 pt-4 dark:border-dark-600" + > + <div> + <label class="input-label mb-0">{{ t('admin.accounts.upstreamBilling.autoProbe') }}</label> + <p class="mt-1 text-xs text-gray-500 dark:text-gray-400"> + {{ t('admin.accounts.upstreamBilling.autoProbeHint') }} + </p> + </div> + <Toggle + v-model="upstreamBillingAutoProbeEnabled" + data-testid="upstream-billing-auto-probe" + :aria-label="t('admin.accounts.upstreamBilling.autoProbe')" + /> + </div> + <!-- Gemini API Key tier selection --> <div v-if="form.platform === 'gemini'"> <label class="input-label">{{ t('admin.accounts.gemini.tier.label') }}</label> @@ -3525,6 +3542,7 @@ import ProxyAdBanner from '@/components/common/ProxyAdBanner.vue' import GroupSelector from '@/components/common/GroupSelector.vue' import ModelWhitelistSelector from '@/components/account/ModelWhitelistSelector.vue' import QuotaLimitCard from '@/components/account/QuotaLimitCard.vue' +import Toggle from '@/components/common/Toggle.vue' import GrokBaseUrlPresets from '@/components/account/GrokBaseUrlPresets.vue' import HeaderOverrideEditor from '@/components/account/HeaderOverrideEditor.vue' import { @@ -3668,6 +3686,7 @@ const accountCategory = ref<'oauth-based' | 'apikey' | 'bedrock' | 'service_acco const addMethod = ref<AddMethod>('oauth') // For oauth-based: 'oauth' or 'setup-token' const apiKeyBaseUrl = ref('https://api.anthropic.com') const apiKeyValue = ref('') +const upstreamBillingAutoProbeEnabled = ref(true) const syncPreviewCredentials = computed(() => { if (!apiKeyValue.value) return undefined @@ -4553,7 +4572,18 @@ const ensureAntigravityMixedChannelConfirmed = async (onConfirm: () => Promise<v const submitCreateAccount = async (payload: CreateAccountRequest) => { submitting.value = true try { - await adminAPI.accounts.create(withAntigravityConfirmFlag(payload)) + const account = await adminAPI.accounts.create(withAntigravityConfirmFlag(payload)) + if ( + payload.platform === 'openai' && + payload.type === 'apikey' && + payload.upstream_billing_probe_enabled === true + ) { + try { + await adminAPI.accounts.probeUpstreamBilling(account.id) + } catch { + appStore.showWarning(t('admin.accounts.upstreamBilling.probeFailed')) + } + } appStore.showSuccess(t('admin.accounts.accountCreated')) emit('created') handleClose() @@ -4593,6 +4623,7 @@ const resetForm = () => { addMethod.value = 'oauth' apiKeyBaseUrl.value = 'https://api.anthropic.com' apiKeyValue.value = '' + upstreamBillingAutoProbeEnabled.value = true editQuotaLimit.value = null editQuotaDailyLimit.value = null editQuotaWeeklyLimit.value = null @@ -5092,6 +5123,8 @@ const handleSubmit = async () => { ...form, group_ids: form.group_ids, extra, + upstream_billing_probe_enabled: + form.platform === 'openai' ? upstreamBillingAutoProbeEnabled.value : undefined, auto_pause_on_expired: autoPauseOnExpired.value }) } diff --git a/frontend/src/components/account/UpstreamBillingRateCell.vue b/frontend/src/components/account/UpstreamBillingRateCell.vue index 73db04d85..98e5671b6 100644 --- a/frontend/src/components/account/UpstreamBillingRateCell.vue +++ b/frontend/src/components/account/UpstreamBillingRateCell.vue @@ -1,37 +1,71 @@ <template> <div v-if="eligible" class="flex h-6 min-w-[7rem] items-center gap-1"> - <HelpTooltip class="-ml-1" width-class="w-64" data-testid="upstream-billing-details"> + <HelpTooltip class="-ml-1" width-class="w-max max-w-[calc(100vw-2rem)]" data-testid="upstream-billing-details"> <template #trigger> <span - class="cursor-help border-b border-dotted border-gray-300 font-mono text-sm font-medium text-gray-800 dark:border-gray-600 dark:text-gray-200" + class="cursor-help border-b border-dotted border-gray-300 text-sm font-medium dark:border-gray-600" + :class="hasEffectiveRate ? 'font-mono text-gray-800 dark:text-gray-200' : statusClass || 'text-gray-400 dark:text-gray-500'" data-testid="upstream-billing-rate" > - {{ effectiveRate }} + {{ primaryValue }} </span> </template> - <div v-if="data" class="space-y-1"> - <p>{{ t('admin.accounts.upstreamBilling.groupRate', { value: data.group_rate_multiplier }) }}</p> - <p v-if="data.user_rate_multiplier != null"> - {{ t('admin.accounts.upstreamBilling.userRate', { value: data.user_rate_multiplier }) }} + <div class="space-y-1"> + <template v-if="hasEffectiveRate && data"> + <p>{{ t('admin.accounts.upstreamBilling.groupRate', { value: data.group_rate_multiplier }) }}</p> + <p v-if="data.user_rate_multiplier != null"> + {{ t('admin.accounts.upstreamBilling.userRate', { value: data.user_rate_multiplier }) }} + </p> + <p> + {{ + 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') + }} + </p> + <p>{{ t('admin.accounts.upstreamBilling.effectiveRate', { value: currentEffectiveRate ?? '-' }) }}</p> + <p>{{ t('admin.accounts.upstreamBilling.updatedAt', { value: formatDate(snapshot?.received_at) }) }}</p> + </template> + <template v-else-if="stale && lastDetectedRate != null"> + <p data-testid="upstream-billing-last-rate"> + {{ t('admin.accounts.upstreamBilling.lastDetectedRate', { value: lastDetectedRate }) }} + </p> + <p data-testid="upstream-billing-last-time"> + {{ t('admin.accounts.upstreamBilling.lastDetectedAt', { value: formatDate(snapshot?.received_at) }) }} + </p> + <p data-testid="upstream-billing-elapsed"> + {{ t('admin.accounts.upstreamBilling.elapsedSince', { value: elapsedSinceLastSuccess }) }} + </p> + </template> + <p v-else>{{ statusLabel || '-' }}</p> + <p + v-if="probeEnabled && globalProbeEnabled !== false && nextProbeAt" + data-testid="upstream-billing-next-probe" + > + {{ t('admin.accounts.upstreamBilling.nextProbeAt', { value: formatDate(nextProbeAt) }) }} </p> - <p> - {{ - 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') - }} + <p class="mt-2 border-t border-white/15 pt-2" data-testid="upstream-billing-probe-state"> + {{ t('admin.accounts.upstreamBilling.accountProbeState') }} + <span :class="probeEnabled ? 'text-emerald-400' : 'text-red-400'"> + {{ probeEnabled ? t('admin.accounts.upstreamBilling.enabled') : t('admin.accounts.upstreamBilling.disabled') }} + </span> + </p> + <p + v-if="globalProbeEnabled === false" + class="mt-1" + data-testid="upstream-billing-global-probe-state" + > + {{ t('admin.accounts.upstreamBilling.globalProbeState') }} + <span class="text-red-400">{{ t('admin.accounts.upstreamBilling.disabled') }}</span> </p> - <p>{{ t('admin.accounts.upstreamBilling.effectiveRate', { value: currentEffectiveRate ?? '-' }) }}</p> - <p>{{ t('admin.accounts.upstreamBilling.updatedAt', { value: formatDate(snapshot?.received_at) }) }}</p> </div> - <span v-else>{{ statusLabel }}</span> </HelpTooltip> - <span v-if="statusLabel" :class="statusClass" class="whitespace-nowrap text-[10px] font-medium"> + <span v-if="hasEffectiveRate && statusLabel" :class="statusClass" class="whitespace-nowrap text-[10px] font-medium"> {{ statusLabel }} </span> <button @@ -56,21 +90,29 @@ import HelpTooltip from '@/components/common/HelpTooltip.vue' import Icon from '@/components/icons/Icon.vue' import type { Account, UpstreamBillingProbeSnapshot } from '@/types' -const props = defineProps<{ +const props = withDefaults(defineProps<{ account: Account - intervalMinutes: number now: number probing?: boolean -}>() + globalProbeEnabled?: boolean +}>(), { + globalProbeEnabled: true +}) defineEmits<{ (event: 'probe'): void }>() const { t } = useI18n() +const CLOCK_SKEW_TOLERANCE_MS = 5 * 60 * 1000 const eligible = computed(() => props.account.platform === 'openai' && props.account.type === 'apikey') const snapshot = computed<UpstreamBillingProbeSnapshot | undefined>(() => props.account.extra?.upstream_billing_probe) const data = computed(() => snapshot.value?.data) +const probeEnabled = computed(() => props.account.extra?.upstream_billing_probe_enabled === true) +const nextProbeAt = computed(() => { + const value = snapshot.value?.next_probe_at + return typeof value === 'string' && Number.isFinite(Date.parse(value)) ? value : '' +}) const receivedAt = computed(() => typeof snapshot.value?.received_at === 'string' ? Date.parse(snapshot.value.received_at) : Number.NaN) const freshUntil = computed(() => { if (typeof snapshot.value?.fresh_until === 'string') return Date.parse(snapshot.value.fresh_until) @@ -81,7 +123,7 @@ const freshUntil = computed(() => { : Number.NaN }) const validTimestamps = computed(() => { - if (!Number.isFinite(receivedAt.value) || receivedAt.value > props.now) return false + if (!Number.isFinite(receivedAt.value) || receivedAt.value > props.now + CLOCK_SKEW_TOLERANCE_MS) return false return Number.isFinite(freshUntil.value) && freshUntil.value > receivedAt.value }) const stale = computed(() => { @@ -130,6 +172,21 @@ const currentEffectiveRate = computed(() => { const value = minute >= start && minute < end ? base * peak : base return Number.isFinite(value) ? value : null }) +const lastDetectedRate = computed(() => { + const value = data.value?.effective_rate_multiplier + return typeof value === 'number' && Number.isFinite(value) && value >= 0 + ? Number(value.toPrecision(12)) + : null +}) +const elapsedSinceLastSuccess = computed(() => { + if (!Number.isFinite(receivedAt.value)) return '-' + const elapsedMinutes = Math.max(0, Math.floor((props.now - receivedAt.value) / 60_000)) + if (elapsedMinutes < 1) return t('admin.accounts.upstreamBilling.justNow') + if (elapsedMinutes < 60) return t('admin.accounts.upstreamBilling.minutesAgo', { count: elapsedMinutes }) + const elapsedHours = Math.floor(elapsedMinutes / 60) + if (elapsedHours < 24) return t('admin.accounts.upstreamBilling.hoursAgo', { count: elapsedHours }) + return t('admin.accounts.upstreamBilling.daysAgo', { count: Math.floor(elapsedHours / 24) }) +}) const effectiveRate = computed(() => { if (!validTimestamps.value || stale.value || !['ok', 'failed'].includes(snapshot.value?.status ?? '')) return '-' const value = currentEffectiveRate.value @@ -149,6 +206,8 @@ const statusClass = computed(() => { if (snapshot.value.status === 'failed') return 'text-red-600 dark:text-red-400' return '' }) +const hasEffectiveRate = computed(() => effectiveRate.value !== '-') +const primaryValue = computed(() => hasEffectiveRate.value ? effectiveRate.value : statusLabel.value || '-') const formatDate = (value?: string) => value ? new Date(value).toLocaleString(undefined, { month: '2-digit', diff --git a/frontend/src/components/account/__tests__/BulkEditAccountModal.spec.ts b/frontend/src/components/account/__tests__/BulkEditAccountModal.spec.ts index 489b00632..cedabc58d 100644 --- a/frontend/src/components/account/__tests__/BulkEditAccountModal.spec.ts +++ b/frontend/src/components/account/__tests__/BulkEditAccountModal.spec.ts @@ -343,6 +343,73 @@ describe('BulkEditAccountModal', () => { }) }) + it('OpenAI API Key 批量编辑可统一开启上游倍率自动探测', async () => { + const wrapper = mountModal({ + selectedPlatforms: ['openai'], + selectedTypes: ['apikey'] + }) + + await wrapper.get('#bulk-edit-upstream-billing-auto-probe-enabled').setValue(true) + await wrapper.get('#bulk-edit-account-form').trigger('submit.prevent') + await flushPromises() + + expect(adminAPI.accounts.bulkUpdate).toHaveBeenCalledTimes(1) + expect(adminAPI.accounts.bulkUpdate).toHaveBeenCalledWith([1, 2], { + upstream_billing_probe_enabled: true + }) + }) + + it('OpenAI API Key 批量编辑可统一关闭上游倍率自动探测', async () => { + const wrapper = mountModal({ + selectedPlatforms: ['openai'], + selectedTypes: ['apikey'] + }) + + await wrapper.get('#bulk-edit-upstream-billing-auto-probe-enabled').setValue(true) + await wrapper.get('[data-testid="bulk-edit-upstream-billing-auto-probe-select"]').setValue('disabled') + await wrapper.get('#bulk-edit-account-form').trigger('submit.prevent') + await flushPromises() + + expect(adminAPI.accounts.bulkUpdate).toHaveBeenCalledTimes(1) + expect(adminAPI.accounts.bulkUpdate).toHaveBeenCalledWith([1, 2], { + upstream_billing_probe_enabled: false + }) + }) + + it('非 OpenAI API Key 目标不显示上游倍率自动探测批量开关', () => { + const wrapper = mountModal({ + selectedPlatforms: ['openai'], + selectedTypes: ['oauth'] + }) + + expect(wrapper.find('#bulk-edit-upstream-billing-auto-probe-enabled').exists()).toBe(false) + }) + + it('筛选结果批量编辑可统一开启上游倍率自动探测', async () => { + const wrapper = mountModal({ + accountIds: [], + selectedPlatforms: [], + selectedTypes: [], + target: { + mode: 'filtered', + filters: { platform: 'openai', type: 'apikey', status: 'active' }, + previewCount: 20, + selectedPlatforms: ['openai'], + selectedTypes: ['apikey'] + } + }) + + await wrapper.get('#bulk-edit-upstream-billing-auto-probe-enabled').setValue(true) + await wrapper.get('#bulk-edit-account-form').trigger('submit.prevent') + await flushPromises() + + expect(adminAPI.accounts.bulkUpdate).toHaveBeenCalledTimes(1) + expect(adminAPI.accounts.bulkUpdate).toHaveBeenCalledWith({ + filters: { platform: 'openai', type: 'apikey', status: 'active' }, + upstream_billing_probe_enabled: true + }) + }) + it('筛选 OpenAI 账号批量编辑应提交 Compact 模式和专属模型映射', async () => { const wrapper = mountModal({ accountIds: [], diff --git a/frontend/src/components/account/__tests__/CreateAccountModal.spec.ts b/frontend/src/components/account/__tests__/CreateAccountModal.spec.ts index 22cee856b..13035a9d3 100644 --- a/frontend/src/components/account/__tests__/CreateAccountModal.spec.ts +++ b/frontend/src/components/account/__tests__/CreateAccountModal.spec.ts @@ -4,10 +4,12 @@ import { beforeEach, describe, expect, it, vi } from 'vitest' const { createAccountMock, + probeUpstreamBillingMock, importCodexSessionMock, createOpenAICodexPATMock, } = vi.hoisted(() => ({ createAccountMock: vi.fn(), + probeUpstreamBillingMock: vi.fn(), importCodexSessionMock: vi.fn(), createOpenAICodexPATMock: vi.fn(), })) @@ -28,6 +30,7 @@ vi.mock('@/api/admin', () => ({ adminAPI: { accounts: { create: createAccountMock, + probeUpstreamBilling: probeUpstreamBillingMock, checkMixedChannelRisk: vi.fn().mockResolvedValue({ has_risk: false }), importCodexSession: importCodexSessionMock, createOpenAICodexPAT: createOpenAICodexPATMock, @@ -108,7 +111,11 @@ async function selectButtonByText(wrapper: ReturnType<typeof mountModal>, text: await button?.trigger('click') } -async function submitApiKeyAccount(platform: 'openai' | 'anthropic', enableLongContextBilling = false) { +async function submitApiKeyAccount( + platform: 'openai' | 'anthropic', + enableLongContextBilling = false, + disableUpstreamBillingProbe = false +) { const wrapper = mountModal() await selectButtonByText(wrapper, platform === 'openai' ? 'OpenAI' : 'admin.accounts.claudeConsole') if (platform === 'openai') { @@ -119,8 +126,12 @@ async function submitApiKeyAccount(platform: 'openai' | 'anthropic', enableLongC if (enableLongContextBilling) { await wrapper.get('[data-testid="openai-long-context-billing-toggle"]').trigger('click') } + if (disableUpstreamBillingProbe) { + await wrapper.get('[data-testid="upstream-billing-auto-probe"]').trigger('click') + } await wrapper.get('form#create-account-form').trigger('submit.prevent') await flushPromises() + return wrapper } async function openCodexImportStep(toggleClicks = 0) { @@ -136,7 +147,8 @@ async function openCodexImportStep(toggleClicks = 0) { describe('CreateAccountModal OpenAI long-context billing', () => { beforeEach(() => { - createAccountMock.mockReset().mockResolvedValue({}) + createAccountMock.mockReset().mockResolvedValue({ id: 42, platform: 'openai', type: 'apikey' }) + probeUpstreamBillingMock.mockReset().mockResolvedValue({}) importCodexSessionMock.mockReset().mockResolvedValue({ created: 1, updated: 0, @@ -155,6 +167,38 @@ describe('CreateAccountModal OpenAI long-context billing', () => { expect(createAccountMock.mock.calls[0]?.[0]?.extra?.openai_long_context_billing_enabled).toBe(false) }) + it('enables upstream billing probes by default for new OpenAI API key accounts', async () => { + await submitApiKeyAccount('openai') + + expect(createAccountMock.mock.calls[0]?.[0]?.upstream_billing_probe_enabled).toBe(true) + }) + + it('waits for the initial upstream billing probe before refreshing the account list', async () => { + let resolveProbe: (() => void) | undefined + probeUpstreamBillingMock.mockImplementationOnce( + () => new Promise<void>((resolve) => { + resolveProbe = resolve + }) + ) + + const wrapper = await submitApiKeyAccount('openai') + + expect(probeUpstreamBillingMock).toHaveBeenCalledWith(42) + expect(wrapper.emitted('created')).toBeUndefined() + + resolveProbe?.() + await flushPromises() + + expect(wrapper.emitted('created')).toHaveLength(1) + }) + + it('sends an explicit disabled state when the create toggle is turned off', async () => { + await submitApiKeyAccount('openai', false, true) + + expect(createAccountMock.mock.calls[0]?.[0]?.upstream_billing_probe_enabled).toBe(false) + expect(probeUpstreamBillingMock).not.toHaveBeenCalled() + }) + it('exposes Agent Identity in the OpenAI authorization methods', async () => { const wrapper = mountModal() await selectButtonByText(wrapper, 'OpenAI') @@ -195,6 +239,7 @@ describe('CreateAccountModal OpenAI long-context billing', () => { expect(createAccountMock).toHaveBeenCalledTimes(1) expect(createAccountMock.mock.calls[0]?.[0]?.extra?.openai_long_context_billing_enabled).toBeUndefined() + expect(createAccountMock.mock.calls[0]?.[0]?.upstream_billing_probe_enabled).toBeUndefined() }) it('leaves Codex session import billing ownership to the backend', async () => { diff --git a/frontend/src/components/account/__tests__/UpstreamBillingRateCell.spec.ts b/frontend/src/components/account/__tests__/UpstreamBillingRateCell.spec.ts index 8246c7066..13fb70478 100644 --- a/frontend/src/components/account/__tests__/UpstreamBillingRateCell.spec.ts +++ b/frontend/src/components/account/__tests__/UpstreamBillingRateCell.spec.ts @@ -1,6 +1,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' -import { mount } from '@vue/test-utils' +import { flushPromises, mount } from '@vue/test-utils' import UpstreamBillingRateCell from '../UpstreamBillingRateCell.vue' +import HelpTooltip from '@/components/common/HelpTooltip.vue' import type { Account } from '@/types' vi.mock('vue-i18n', async () => { @@ -83,7 +84,6 @@ describe('UpstreamBillingRateCell', () => { } } }), - intervalMinutes: 30, now: Date.now() } }) @@ -113,9 +113,9 @@ describe('UpstreamBillingRateCell', () => { } } }) - const wrapper = mount(UpstreamBillingRateCell, { props: { account, intervalMinutes: 30, now: Date.now() } }) + const wrapper = mount(UpstreamBillingRateCell, { props: { account, now: Date.now() } }) expect(wrapper.text()).toContain('admin.accounts.upstreamBilling.stale') - expect(wrapper.get('[data-testid="upstream-billing-rate"]').text()).toBe('-') + expect(wrapper.get('[data-testid="upstream-billing-rate"]').text()).toBe('admin.accounts.upstreamBilling.stale') await wrapper.setProps({ account: makeAccount({ @@ -140,7 +140,7 @@ describe('UpstreamBillingRateCell', () => { expect(wrapper.text()).not.toContain('admin.accounts.upstreamBilling.stale') await wrapper.setProps({ now: Date.parse('2026-07-13T01:00:00.001Z') }) - expect(wrapper.get('[data-testid="upstream-billing-rate"]').text()).toBe('-') + expect(wrapper.get('[data-testid="upstream-billing-rate"]').text()).toBe('admin.accounts.upstreamBilling.stale') expect(wrapper.text()).toContain('admin.accounts.upstreamBilling.stale') await wrapper.setProps({ @@ -159,13 +159,111 @@ describe('UpstreamBillingRateCell', () => { } }) }) - expect(wrapper.get('[data-testid="upstream-billing-rate"]').text()).toBe('-') + expect(wrapper.get('[data-testid="upstream-billing-rate"]').text()).toBe('admin.accounts.upstreamBilling.stale') expect(wrapper.text()).toContain('admin.accounts.upstreamBilling.stale') }) + it('shows stale snapshot details, local next probe time, and the account probe state', async () => { + const wrapper = mount(UpstreamBillingRateCell, { + attachTo: document.body, + props: { + account: makeAccount({ + extra: { + upstream_billing_probe_enabled: true, + upstream_billing_probe: { + status: 'ok', + data: billingData, + received_at: '2026-07-12T22:00:00Z', + fresh_until: '2026-07-12T23:00:00Z', + last_attempt_at: '2026-07-12T22:00:00Z', + next_probe_at: '2026-07-13T01:00:00Z' + } + } + }), + now: Date.now() + } + }) + + expect(wrapper.getComponent(HelpTooltip).props('widthClass')).toBe('w-max max-w-[calc(100vw-2rem)]') + await wrapper.get('[data-testid="upstream-billing-details"]').trigger('mouseenter') + await flushPromises() + + const tooltips = document.body.querySelectorAll('[role="tooltip"]') + const tooltip = tooltips[tooltips.length - 1] as HTMLElement + expect(tooltip.textContent).toContain('admin.accounts.upstreamBilling.lastDetectedRate:0.9') + expect(tooltip.textContent).toContain('admin.accounts.upstreamBilling.lastDetectedAt:') + expect(tooltip.textContent).toContain('admin.accounts.upstreamBilling.elapsedSince:admin.accounts.upstreamBilling.hoursAgo:2') + expect(tooltip.textContent).toContain('admin.accounts.upstreamBilling.nextProbeAt:') + expect(tooltip.textContent).not.toContain('admin.accounts.upstreamBilling.stale') + expect(tooltip.querySelector('[data-testid="upstream-billing-probe-state"] span')?.className).toContain('text-emerald-400') + + await wrapper.setProps({ + account: makeAccount({ + extra: { + upstream_billing_probe_enabled: false, + upstream_billing_probe: { + status: 'unsupported', + last_attempt_at: '2026-07-13T00:00:00Z', + next_probe_at: '2026-07-13T01:00:00Z' + } + } + }) + }) + expect(tooltip.querySelector('[data-testid="upstream-billing-next-probe"]')).toBeNull() + expect(tooltip.querySelector('[data-testid="upstream-billing-probe-state"] span')?.className).toContain('text-red-400') + wrapper.unmount() + }) + + it('stacks the global-off state below the account state and hides it when globally enabled', async () => { + const wrapper = mount(UpstreamBillingRateCell, { + attachTo: document.body, + props: { + account: makeAccount({ + extra: { + upstream_billing_probe_enabled: true, + upstream_billing_probe: { + status: 'unsupported', + last_attempt_at: '2026-07-13T00:00:00Z', + next_probe_at: '2026-07-13T01:00:00Z' + } + } + }), + globalProbeEnabled: false, + now: Date.now() + } + }) + + await wrapper.get('[data-testid="upstream-billing-details"]').trigger('mouseenter') + await flushPromises() + + const tooltips = document.body.querySelectorAll('[role="tooltip"]') + const tooltip = tooltips[tooltips.length - 1] as HTMLElement + const accountState = tooltip.querySelector('[data-testid="upstream-billing-probe-state"]') + const globalState = tooltip.querySelector('[data-testid="upstream-billing-global-probe-state"]') + expect(accountState?.querySelector('span')?.className).toContain('text-emerald-400') + expect(globalState?.textContent).toContain('admin.accounts.upstreamBilling.globalProbeState') + expect(globalState?.querySelector('span')?.className).toContain('text-red-400') + expect(tooltip.querySelector('[data-testid="upstream-billing-next-probe"]')).toBeNull() + + await wrapper.setProps({ globalProbeEnabled: true }) + expect(tooltip.querySelector('[data-testid="upstream-billing-global-probe-state"]')).toBeNull() + expect(tooltip.querySelector('[data-testid="upstream-billing-next-probe"]')).not.toBeNull() + + await wrapper.setProps({ + globalProbeEnabled: false, + account: makeAccount({ + extra: { upstream_billing_probe_enabled: false } + }) + }) + expect(accountState?.querySelector('span')?.className).toContain('text-red-400') + expect(tooltip.querySelector('[data-testid="upstream-billing-global-probe-state"]')).not.toBeNull() + expect(tooltip.querySelector('[data-testid="upstream-billing-next-probe"]')).toBeNull() + wrapper.unmount() + }) + it('emits manual probe commands only for eligible accounts', async () => { const wrapper = mount(UpstreamBillingRateCell, { - props: { account: makeAccount(), intervalMinutes: 30, now: Date.now() } + props: { account: makeAccount(), now: Date.now() } }) await wrapper.get('[data-testid="upstream-billing-probe"]').trigger('click') expect(wrapper.emitted('probe')).toHaveLength(1) @@ -199,7 +297,6 @@ describe('UpstreamBillingRateCell', () => { peak_rate_enabled: false, effective_rate_multiplier: -1 }), - intervalMinutes: 30, now: Date.now() } }) @@ -208,11 +305,13 @@ describe('UpstreamBillingRateCell', () => { await wrapper.setProps({ account: malformedAccount({ billing_scope: 'request' as 'token' }) }) expect(wrapper.get('[data-testid="upstream-billing-rate"]').text()).toBe('-') await wrapper.setProps({ account: malformedAccount({}, { received_at: 'not-a-time' }) }) - expect(wrapper.get('[data-testid="upstream-billing-rate"]').text()).toBe('-') + expect(wrapper.get('[data-testid="upstream-billing-rate"]').text()).toBe('admin.accounts.upstreamBilling.stale') await wrapper.setProps({ account: malformedAccount({}, { received_at: '2026-07-13T00:31:00Z' }) }) - expect(wrapper.get('[data-testid="upstream-billing-rate"]').text()).toBe('-') + expect(wrapper.get('[data-testid="upstream-billing-rate"]').text()).toBe('0.6x') + await wrapper.setProps({ account: malformedAccount({}, { received_at: '2026-07-13T00:36:00Z' }) }) + expect(wrapper.get('[data-testid="upstream-billing-rate"]').text()).toBe('admin.accounts.upstreamBilling.stale') await wrapper.setProps({ account: malformedAccount({}, { fresh_until: '2026-07-12T23:59:00Z' }) }) - expect(wrapper.get('[data-testid="upstream-billing-rate"]').text()).toBe('-') + expect(wrapper.get('[data-testid="upstream-billing-rate"]').text()).toBe('admin.accounts.upstreamBilling.stale') await wrapper.setProps({ account: makeAccount({ @@ -226,8 +325,31 @@ describe('UpstreamBillingRateCell', () => { } }) }) - expect(wrapper.get('[data-testid="upstream-billing-rate"]').text()).toBe('-') + expect(wrapper.get('[data-testid="upstream-billing-rate"]').text()).toBe('admin.accounts.upstreamBilling.failed') expect(wrapper.text()).toContain('admin.accounts.upstreamBilling.failed') expect(wrapper.text()).not.toContain('admin.accounts.upstreamBilling.stale') }) + + it('uses unsupported as the primary tooltip trigger without a dash', () => { + const wrapper = mount(UpstreamBillingRateCell, { + props: { + account: makeAccount({ + extra: { + upstream_billing_probe: { + status: 'unsupported', + last_attempt_at: '2026-07-13T00:00:00Z', + next_probe_at: '2026-07-13T00:30:00Z', + last_error: 'unsupported' + } + } + }), + now: Date.now() + } + }) + + expect(wrapper.get('[data-testid="upstream-billing-rate"]').text()).toBe( + 'admin.accounts.upstreamBilling.unsupported' + ) + expect(wrapper.text()).not.toContain('-admin.accounts.upstreamBilling.unsupported') + }) }) diff --git a/frontend/src/components/common/DataTable.vue b/frontend/src/components/common/DataTable.vue index 9dcca93c9..d593f8f9a 100644 --- a/frontend/src/components/common/DataTable.vue +++ b/frontend/src/components/common/DataTable.vue @@ -130,38 +130,38 @@ ]" @click="column.sortable && handleSort(column.key)" > - <slot - :name="`header-${column.key}`" - :column="column" - :sort-key="sortKey" - :sort-order="sortOrder" - > - <div :class="['flex items-center space-x-1', getHeaderContentAlignmentClass(column)]"> + <div :class="['flex items-center space-x-1', getHeaderContentAlignmentClass(column)]"> + <slot + :name="`header-${column.key}`" + :column="column" + :sort-key="sortKey" + :sort-order="sortOrder" + > <span>{{ column.label }}</span> - <span - v-if="column.sortable" - class="inline-flex h-5 w-4 flex-col items-center justify-center" - aria-hidden="true" + </slot> + <span + v-if="column.sortable" + class="inline-flex h-5 w-4 flex-col items-center justify-center" + aria-hidden="true" + > + <svg + class="h-2.5 w-2.5" + :class="getSortIndicatorClass(column.key, 'asc')" + fill="currentColor" + viewBox="0 0 10 10" > - <svg - class="h-2.5 w-2.5" - :class="getSortIndicatorClass(column.key, 'asc')" - fill="currentColor" - viewBox="0 0 10 10" - > - <path d="M5 2L1.5 6.5h7L5 2z" /> - </svg> - <svg - class="-mt-0.5 h-2.5 w-2.5" - :class="getSortIndicatorClass(column.key, 'desc')" - fill="currentColor" - viewBox="0 0 10 10" - > - <path d="M5 8L1.5 3.5h7L5 8z" /> - </svg> - </span> - </div> - </slot> + <path d="M5 2L1.5 6.5h7L5 2z" /> + </svg> + <svg + class="-mt-0.5 h-2.5 w-2.5" + :class="getSortIndicatorClass(column.key, 'desc')" + fill="currentColor" + viewBox="0 0 10 10" + > + <path d="M5 8L1.5 3.5h7L5 8z" /> + </svg> + </span> + </div> </th> </tr> </thead> diff --git a/frontend/src/components/common/__tests__/DataTable.spec.ts b/frontend/src/components/common/__tests__/DataTable.spec.ts index de1407c8c..91bbbc924 100644 --- a/frontend/src/components/common/__tests__/DataTable.spec.ts +++ b/frontend/src/components/common/__tests__/DataTable.spec.ts @@ -60,12 +60,16 @@ describe('DataTable', () => { ], defaultSortKey: 'name', defaultSortOrder: 'asc' + }, + slots: { + 'header-name': '<span data-test="custom-name-header">Name</span>' } }) await wrapper.vm.$nextTick() const nameHeader = wrapper.findAll('th')[0] + expect(nameHeader.find('[data-test="custom-name-header"]').exists()).toBe(true) expect(nameHeader.attributes('aria-sort')).toBe('ascending') expect(nameHeader.findAll('svg')).toHaveLength(2) expect(nameHeader.findAll('svg')[0].classes()).toContain('text-primary-600') diff --git a/frontend/src/i18n/__tests__/wsModeLocaleDesc.spec.ts b/frontend/src/i18n/__tests__/wsModeLocaleDesc.spec.ts new file mode 100644 index 000000000..43075e9d8 --- /dev/null +++ b/frontend/src/i18n/__tests__/wsModeLocaleDesc.spec.ts @@ -0,0 +1,13 @@ +import { describe, expect, it } from 'vitest' + +import en from '../locales/en/admin/accounts' +import zh from '../locales/zh/admin/accounts' + +describe('OpenAI WS mode locale descriptions', () => { + it('documents the global v2 router requirement for account WS modes', () => { + expect(zh.accounts.openai.wsModeDesc).toContain('mode_router_v2_enabled') + expect(zh.accounts.openai.wsModeDesc).toContain('http_bridge') + expect(en.accounts.openai.wsModeDesc).toContain('mode_router_v2_enabled') + expect(en.accounts.openai.wsModeDesc).toContain('http_bridge') + }) +}) diff --git a/frontend/src/i18n/locales/en/admin/accounts.ts b/frontend/src/i18n/locales/en/admin/accounts.ts index f60dca15e..206c7c517 100644 --- a/frontend/src/i18n/locales/en/admin/accounts.ts +++ b/frontend/src/i18n/locales/en/admin/accounts.ts @@ -175,10 +175,8 @@ export default { usageWindowsHint: '"5h / 7d" are the upstream account\'s official rolling usage windows (e.g. OpenAI ChatGPT, Claude). They are imposed by the upstream provider on the account itself — not configured by sub2api, and unrelated to the models you map. Usage resets automatically once each window rolls over, and the limit cannot be lifted from within sub2api.', upstreamBilling: { trustWarning: 'This rate is declared by the upstream site for the current API key. Sub2API cannot verify that it matches actual charges. The upstream site or an intermediary may return forged, stale, or modified data. Verify it against bills, balance changes, and actual usage.', - autoProbeSettings: 'Upstream rate auto probe', - intervalMinutes: 'Probe interval (minutes)', - autoProbe: 'Auto probe', - autoProbeHint: 'Probe this account on the global interval when global probing is enabled.', + autoProbe: 'Automatically probe upstream declared rate', + autoProbeHint: 'Probe this account\'s upstream declared rate on the global interval when global probing is enabled.', manualProbe: 'Probe upstream rate now', stale: 'Stale', unsupported: 'Unsupported', @@ -190,8 +188,18 @@ export default { noPeakRate: 'Peak rate: disabled', effectiveRate: 'Current rate: {value}x', updatedAt: 'Updated: {value}', - settingsSaved: 'Upstream rate probe settings saved', - settingsFailed: 'Failed to save upstream rate probe settings', + nextProbeAt: 'Next probe: {value}', + lastDetectedRate: 'Last detected rate: {value}x', + lastDetectedAt: 'Last detected: {value}', + elapsedSince: 'Elapsed: {value}', + justNow: 'less than 1 minute', + minutesAgo: '{count} minutes', + hoursAgo: '{count} hours', + daysAgo: '{count} days', + accountProbeState: 'Automatic detection for this account:', + globalProbeState: 'Global probe switch:', + enabled: 'On', + disabled: 'Off', probeFailed: 'Failed to probe upstream rate', noEligibleAccounts: 'Select OpenAI API key accounts', batchLimit: 'A batch can probe at most 20 accounts', @@ -439,7 +447,8 @@ export default { responsesWebsocketsV2Desc: 'Disabled by default. Enable to allow responses_websockets_v2 capability (still gated by global and account-type switches).', wsMode: 'WS mode', - wsModeDesc: 'Only applies to the current OpenAI account type.', + wsModeDesc: + 'Only applies to the current OpenAI account type; account WS modes, including http_bridge, take effect only when the global gateway.openai_ws.mode_router_v2_enabled=true.', wsModeOff: 'Off (off)', wsModeCtxPool: 'Context Pool (ctx_pool)', wsModePassthrough: 'Passthrough (passthrough)', diff --git a/frontend/src/i18n/locales/en/admin/ops.ts b/frontend/src/i18n/locales/en/admin/ops.ts index 6a6d02a81..b6df3a911 100644 --- a/frontend/src/i18n/locales/en/admin/ops.ts +++ b/frontend/src/i18n/locales/en/admin/ops.ts @@ -381,8 +381,6 @@ export default { suggestPlatform: 'Platform error: prioritize investigation and fix', suggestGeneric: 'See details for more context', apiKeyPrefix: 'Key Prefix', - attemptedKeyPrefix: 'Attempted Key Prefix', - deletedKeyOwner: 'Deleted Key Owner', keyDeletedBadge: 'Key Deleted' }, requestDetails: { @@ -712,8 +710,6 @@ export default { ignoreContextCanceledHint: 'When enabled, client disconnect (context canceled) errors will not be written to the error log.', ignoreNoAvailableAccounts: 'Ignore no available accounts errors', ignoreNoAvailableAccountsHint: 'When enabled, "No available accounts" errors will not be written to the error log (not recommended; usually a config issue).', - ignoreInvalidApiKeyErrors: 'Ignore invalid API key errors', - ignoreInvalidApiKeyErrorsHint: 'When enabled, invalid or missing API key errors (INVALID_API_KEY, API_KEY_REQUIRED) will not be written to the error log.', ignoreInsufficientBalanceErrors: 'Ignore Insufficient Balance Errors', ignoreInsufficientBalanceErrorsHint: 'When enabled, insufficient account balance errors will not be written to the error log.', autoRefresh: 'Auto Refresh', diff --git a/frontend/src/i18n/locales/en/admin/settings.ts b/frontend/src/i18n/locales/en/admin/settings.ts index 8c57ab415..531769eec 100644 --- a/frontend/src/i18n/locales/en/admin/settings.ts +++ b/frontend/src/i18n/locales/en/admin/settings.ts @@ -316,6 +316,16 @@ export default { allowUngroupedKey: 'Allow Ungrouped Key Scheduling', allowUngroupedKeyHint: 'When disabled, API Keys not assigned to any group cannot make requests (403 Forbidden). Keep disabled to ensure all Keys belong to a specific group.' }, + upstreamBillingProbe: { + title: 'Upstream Rate Auto Detection', + description: 'Periodically retrieve billing rates declared by upstream Sub2API sites connected to OpenAI API keys.', + enabled: 'Enable global auto detection', + enabledHint: 'When enabled, scheduled detection runs only for accounts that also enable automatic detection. Disabling stops all scheduled detection; manual detection remains available.', + intervalMinutes: 'Detection interval (minutes)', + intervalHint: 'Range: 5–1440 minutes. A successful result remains valid for two detection intervals.', + saved: 'Upstream rate auto detection settings saved', + saveFailed: 'Failed to save upstream rate auto detection settings' + }, gatewayForwarding: { title: 'Request Forwarding', description: 'Control how requests are forwarded to upstream OAuth accounts', diff --git a/frontend/src/i18n/locales/en/misc.ts b/frontend/src/i18n/locales/en/misc.ts index d0f35b36d..eae9d237d 100644 --- a/frontend/src/i18n/locales/en/misc.ts +++ b/frontend/src/i18n/locales/en/misc.ts @@ -550,7 +550,7 @@ export default { currencyHint: 'Display-only 3-letter ISO currency code shown next to the price; leave empty to hide, does not affect billing', subscriptionCnyPayPreview: 'CNY channel charge preview: {amount}', subscriptionCnyPayPreviewWithFee: '({feeRate}% fee included: {total})', - validityDays: 'Validity (days)', + validity: 'Validity', validityUnit: 'Validity Unit', sortOrder: 'Sort Order', forSale: 'For Sale', @@ -588,7 +588,7 @@ export default { selectGroup: 'Select a group', groupRequired: 'Please select a subscription group', priceRequired: 'Price must be greater than 0', - validityDaysRequired: 'Validity days must be greater than 0', + validityRequired: 'Validity must be greater than 0', groupMissing: 'Missing', groupInfo: 'Group Info', platform: 'Platform', diff --git a/frontend/src/i18n/locales/zh/admin/accounts.ts b/frontend/src/i18n/locales/zh/admin/accounts.ts index 0a6c7392d..48da50175 100644 --- a/frontend/src/i18n/locales/zh/admin/accounts.ts +++ b/frontend/src/i18n/locales/zh/admin/accounts.ts @@ -131,10 +131,8 @@ export default { usageWindowsHint: '“5h / 7d”是上游账号(如 OpenAI ChatGPT、Claude)官方的滚动用量窗口限制,由上游对账号设定,并非 sub2api 配置,也与你映射的模型无关。窗口滚动到期后用量会自动重置,无法在 sub2api 端解除该限制。', upstreamBilling: { trustWarning: '此倍率由上游站点针对当前 API Key 自行声明。Sub2API 无法验证该值是否与实际扣费一致;上游站点或中间代理可能返回伪造、过期或被篡改的数据。请结合账单、余额变化和实际用量自行核验。', - autoProbeSettings: '上游倍率自动探测', - intervalMinutes: '探测周期(分钟)', - autoProbe: '自动探测', - autoProbeHint: '启用后按全局探测周期查询此账号;全局探测关闭时不会执行。', + autoProbe: '自动探测上游声明倍率', + autoProbeHint: '启用后按全局探测周期查询此账号的上游声明倍率;全局探测关闭时不会执行。', manualProbe: '立即探测上游倍率', stale: '已过期', unsupported: '不支持', @@ -146,8 +144,18 @@ export default { noPeakRate: '高峰倍率:未启用', effectiveRate: '当前倍率:{value}x', updatedAt: '更新时间:{value}', - settingsSaved: '上游倍率探测设置已保存', - settingsFailed: '保存上游倍率探测设置失败', + nextProbeAt: '下一次探测:{value}', + lastDetectedRate: '上次探测倍率:{value}x', + lastDetectedAt: '上次探测时间:{value}', + elapsedSince: '已过去:{value}', + justNow: '不足 1 分钟', + minutesAgo: '{count} 分钟', + hoursAgo: '{count} 小时', + daysAgo: '{count} 天', + accountProbeState: '当前账号自动检测:', + globalProbeState: '全局探测开关:', + enabled: '打开', + disabled: '关闭', probeFailed: '探测上游倍率失败', noEligibleAccounts: '请选择 OpenAI API Key 账号', batchLimit: '每次最多探测 20 个账号', @@ -543,7 +551,8 @@ export default { responsesWebsocketsV2Desc: '默认关闭。开启后可启用 responses_websockets_v2 协议能力(受网关全局开关与账号类型开关约束)。', wsMode: 'WS mode', - wsModeDesc: '仅对当前 OpenAI 账号类型生效。', + wsModeDesc: + '仅对当前 OpenAI 账号类型生效;包括 http_bridge 在内的账号 WS mode 仅在全局 gateway.openai_ws.mode_router_v2_enabled=true 时生效。', wsModeOff: '关闭(off)', wsModeCtxPool: '上下文池(ctx_pool)', wsModePassthrough: '透传(passthrough)', diff --git a/frontend/src/i18n/locales/zh/admin/ops.ts b/frontend/src/i18n/locales/zh/admin/ops.ts index 2816ec324..d79dd6ab4 100644 --- a/frontend/src/i18n/locales/zh/admin/ops.ts +++ b/frontend/src/i18n/locales/zh/admin/ops.ts @@ -381,8 +381,6 @@ export default { suggestPlatform: '🚨 平台错误,建议立即排查修复', suggestGeneric: '查看详情了解更多信息', apiKeyPrefix: 'Key 前缀', - attemptedKeyPrefix: '尝试的 Key 前缀', - deletedKeyOwner: '已删除 Key 所有者', keyDeletedBadge: 'Key 已删除' }, requestDetails: { @@ -713,8 +711,6 @@ export default { '启用后,客户端主动断开连接(context canceled)的错误将不会写入错误日志。', ignoreNoAvailableAccounts: '忽略无可用账号错误', ignoreNoAvailableAccountsHint: '启用后,"No available accounts" 错误将不会写入错误日志(不推荐,这通常是配置问题)。', - ignoreInvalidApiKeyErrors: '忽略无效 API Key 错误', - ignoreInvalidApiKeyErrorsHint: '启用后,无效或缺失 API Key 的错误(INVALID_API_KEY、API_KEY_REQUIRED)将不会写入错误日志。', ignoreInsufficientBalanceErrors: '忽略余额不足错误', ignoreInsufficientBalanceErrorsHint: '启用后,账号余额不足(Insufficient balance)的错误将不会写入错误日志。', autoRefresh: '自动刷新', diff --git a/frontend/src/i18n/locales/zh/admin/settings.ts b/frontend/src/i18n/locales/zh/admin/settings.ts index 49f43dc05..f656ac488 100644 --- a/frontend/src/i18n/locales/zh/admin/settings.ts +++ b/frontend/src/i18n/locales/zh/admin/settings.ts @@ -309,6 +309,16 @@ export default { allowUngroupedKey: '允许未分组 Key 调度', allowUngroupedKeyHint: '关闭后,未分配到任何分组的 API Key 将无法发起请求(返回 403)。建议保持关闭以确保所有 Key 都归属明确的分组。' }, + upstreamBillingProbe: { + title: '上游倍率自动探测', + description: '定期获取 OpenAI API Key 所连接上游 Sub2API 站点声明的计费倍率。', + enabled: '启用全局自动探测', + enabledHint: '开启后,仅对账号自身已启用自动检测的账号执行定时探测;关闭后停止所有定时探测,手动探测不受影响。', + intervalMinutes: '探测周期(分钟)', + intervalHint: '范围 5–1440 分钟。成功探测结果的有效期为两个探测周期。', + saved: '上游倍率自动探测设置已保存', + saveFailed: '保存上游倍率自动探测设置失败' + }, gatewayForwarding: { title: '请求转发行为', description: '控制请求转发到上游 OAuth 账号时的行为', diff --git a/frontend/src/i18n/locales/zh/misc.ts b/frontend/src/i18n/locales/zh/misc.ts index f9b975974..31247608d 100644 --- a/frontend/src/i18n/locales/zh/misc.ts +++ b/frontend/src/i18n/locales/zh/misc.ts @@ -574,7 +574,7 @@ export default { currencyHint: '仅用于价格展示的 ISO 三字母币种码,留空不展示,不影响实际扣款', subscriptionCnyPayPreview: 'CNY 通道实扣预览:{amount}', subscriptionCnyPayPreviewWithFee: '(含 {feeRate}% 手续费:{total})', - validityDays: '有效期(天)', + validity: '有效期', validityUnit: '有效期单位', sortOrder: '排序', forSale: '上架状态', @@ -612,7 +612,7 @@ export default { selectGroup: '请选择分组', groupRequired: '请选择订阅分组', priceRequired: '价格必须大于 0', - validityDaysRequired: '有效期天数必须大于 0', + validityRequired: '有效期必须大于 0', groupMissing: '缺失', groupInfo: '分组信息', platform: '平台', diff --git a/frontend/src/main.ts b/frontend/src/main.ts index 68ace8857..d7c17b466 100644 --- a/frontend/src/main.ts +++ b/frontend/src/main.ts @@ -4,6 +4,7 @@ import App from './App.vue' import router from './router' import i18n, { initI18n } from './i18n' import { useAppStore } from '@/stores/app' +import { updateFavicon } from '@/utils/branding' import './style.css' function initThemeClass() { @@ -31,6 +32,7 @@ async function bootstrap() { if (appStore.siteName && appStore.siteName !== 'Sub2API') { document.title = `${appStore.siteName} - AI API Gateway` } + updateFavicon(appStore.siteLogo) await initI18n() diff --git a/frontend/src/types/index.ts b/frontend/src/types/index.ts index 00ebed3fb..5a70bc04b 100644 --- a/frontend/src/types/index.ts +++ b/frontend/src/types/index.ts @@ -1206,6 +1206,7 @@ export interface CreateAccountRequest { group_ids?: number[] expires_at?: number | null auto_pause_on_expired?: boolean + upstream_billing_probe_enabled?: boolean confirm_mixed_channel_risk?: boolean } diff --git a/frontend/src/utils/__tests__/branding.spec.ts b/frontend/src/utils/__tests__/branding.spec.ts new file mode 100644 index 000000000..192c180c1 --- /dev/null +++ b/frontend/src/utils/__tests__/branding.spec.ts @@ -0,0 +1,22 @@ +import { beforeEach, describe, expect, it } from 'vitest' +import { updateFavicon } from '@/utils/branding' + +describe('updateFavicon', () => { + beforeEach(() => { + document.head.innerHTML = '<link rel="icon" href="/logo.png">' + }) + + it('replaces the default favicon with the configured logo', () => { + updateFavicon('https://example.com/custom-logo.png') + + const link = document.querySelector<HTMLLinkElement>('link[rel="icon"]') + expect(link?.href).toBe('https://example.com/custom-logo.png') + }) + + it('ignores unsafe logo URLs', () => { + updateFavicon('javascript:alert(1)') + + const link = document.querySelector<HTMLLinkElement>('link[rel="icon"]') + expect(link?.getAttribute('href')).toBe('/logo.png') + }) +}) diff --git a/frontend/src/utils/branding.ts b/frontend/src/utils/branding.ts new file mode 100644 index 000000000..be8c68c7c --- /dev/null +++ b/frontend/src/utils/branding.ts @@ -0,0 +1,21 @@ +import { sanitizeUrl } from '@/utils/url' + +export function updateFavicon(logoUrl: string): void { + const sanitizedLogoUrl = sanitizeUrl(logoUrl, { + allowRelative: true, + allowDataUrl: true, + }) + if (!sanitizedLogoUrl) { + return + } + + let link = document.querySelector<HTMLLinkElement>('link[rel="icon"]') + if (!link) { + link = document.createElement('link') + link.rel = 'icon' + document.head.appendChild(link) + } + + link.type = sanitizedLogoUrl.endsWith('.svg') ? 'image/svg+xml' : 'image/x-icon' + link.href = sanitizedLogoUrl +} diff --git a/frontend/src/views/admin/AccountsView.vue b/frontend/src/views/admin/AccountsView.vue index fce956e8e..e3bd2266f 100644 --- a/frontend/src/views/admin/AccountsView.vue +++ b/frontend/src/views/admin/AccountsView.vue @@ -132,41 +132,6 @@ <span class="flex-1 text-left">{{ t('admin.tlsFingerprintProfiles.title') }}</span> </button> - <div class="my-2 border-t border-gray-100 dark:border-gray-700"></div> - <div class="space-y-2 px-3 py-2"> - <div class="flex items-center justify-between gap-3"> - <span class="text-sm font-medium text-gray-700 dark:text-gray-200"> - {{ t('admin.accounts.upstreamBilling.autoProbeSettings') }} - </span> - <Toggle - v-model="upstreamBillingProbeSettings.enabled" - :aria-label="t('admin.accounts.upstreamBilling.autoProbeSettings')" - /> - </div> - <div class="flex items-center gap-2"> - <label class="flex-1 text-xs text-gray-500 dark:text-gray-400" for="upstream-billing-probe-interval"> - {{ t('admin.accounts.upstreamBilling.intervalMinutes') }} - </label> - <input - id="upstream-billing-probe-interval" - v-model.number="upstreamBillingProbeSettings.interval_minutes" - type="number" - min="5" - max="1440" - class="input h-8 w-20 px-2 text-sm" - /> - <button - type="button" - class="btn btn-secondary h-8 px-2" - :disabled="upstreamBillingSettingsLoading || upstreamBillingSettingsSaving" - :title="t('common.save')" - @click="saveUpstreamBillingProbeSettings" - > - <Icon name="check" size="sm" /> - </button> - </div> - </div> - <div class="my-2 border-t border-gray-100 dark:border-gray-700"></div> <div class="px-2 py-2"> <div class="flex items-center justify-between gap-3"> @@ -375,15 +340,17 @@ </span> </template> <template #header-upstream_billing_rate="{ column }"> - <div class="flex items-center"> + <div class="flex items-center gap-1"> <span>{{ column.label }}</span> - <HelpTooltip :content="t('admin.accounts.upstreamBilling.trustWarning')" width-class="w-80" /> + <span @click.stop> + <HelpTooltip :content="t('admin.accounts.upstreamBilling.trustWarning')" width-class="w-80" /> + </span> </div> </template> <template #cell-upstream_billing_rate="{ row }"> <UpstreamBillingRateCell :account="row" - :interval-minutes="upstreamBillingProbeSettings.interval_minutes" + :global-probe-enabled="upstreamBillingProbeGloballyEnabled" :now="upstreamBillingNow" :probing="probingUpstreamBilling.has(row.id)" @probe="handleProbeUpstreamBilling(row)" @@ -512,7 +479,6 @@ import AppLayout from '@/components/layout/AppLayout.vue' import TablePageLayout from '@/components/layout/TablePageLayout.vue' import DataTable from '@/components/common/DataTable.vue' import HelpTooltip from '@/components/common/HelpTooltip.vue' -import Toggle from '@/components/common/Toggle.vue' import Pagination from '@/components/common/Pagination.vue' import ConfirmDialog from '@/components/common/ConfirmDialog.vue' import { CreateAccountModal, EditAccountModal, BulkEditAccountModal, SyncFromCrsModal, TempUnschedStatusModal } from '@/components/account' @@ -541,7 +507,7 @@ import { formatDateTime, formatRelativeTime } from '@/utils/format' import { proxyExpiryBadgeClass, proxyExpiryLabelKey } from '@/utils/proxyExpiry' import { extractApiErrorMessage } from '@/utils/apiError' import { sanitizeUrl } from '@/utils/url' -import type { Account, AccountPlatform, AccountSchedulerGroupScore, AccountType, Proxy as AccountProxy, AdminGroup, WindowStats, ClaudeModel, UpstreamBillingProbeSettings, UpstreamBillingProbeSnapshot } from '@/types' +import type { Account, AccountPlatform, AccountSchedulerGroupScore, AccountType, Proxy as AccountProxy, AdminGroup, WindowStats, ClaudeModel, UpstreamBillingProbeSnapshot } from '@/types' const { t } = useI18n() const appStore = useAppStore() @@ -619,14 +585,10 @@ const scheduleModelOptions = ref<SelectOption[]>([]) const togglingSchedulable = ref<number | null>(null) const menu = reactive<{show:boolean, acc:Account|null, pos:{top:number, left:number}|null}>({ show: false, acc: null, pos: null }) const exportingData = ref(false) -const upstreamBillingProbeSettings = reactive<UpstreamBillingProbeSettings>({ - enabled: true, - interval_minutes: 30 -}) -const upstreamBillingSettingsLoading = ref(false) -const upstreamBillingSettingsSaving = ref(false) const probingUpstreamBilling = reactive(new Set<number>()) +const upstreamBillingProbeGloballyEnabled = ref<boolean | undefined>(undefined) const upstreamBillingNow = ref(Date.now()) +let lastUpstreamBillingSortRefreshMinute = -1 useIntervalFn(() => { upstreamBillingNow.value = Date.now() }, 60_000) // Account tools dropdown @@ -653,6 +615,7 @@ const ACCOUNT_SORTABLE_KEYS = new Set([ 'schedulable', 'priority', 'rate_multiplier', + 'upstream_billing_rate', 'last_used_at', 'created_at', 'expires_at' @@ -968,8 +931,15 @@ const resetAutoRefreshCache = () => { const isFirstLoad = ref(true) +function markUpstreamBillingSortRefresh() { + if (sortState.sort_by === 'upstream_billing_rate') { + lastUpstreamBillingSortRefreshMinute = Math.floor(Date.now() / 60_000) + } +} + const load = async () => { const requestParams = params as any + markUpstreamBillingSortRefresh() syncAccountListDerivedParams() hasPendingListSync.value = false resetAutoRefreshCache() @@ -986,6 +956,7 @@ const load = async () => { } const reload = async () => { + markUpstreamBillingSortRefresh() syncAccountListDerivedParams() hasPendingListSync.value = false resetAutoRefreshCache() @@ -994,6 +965,19 @@ const reload = async () => { await refreshTodayStatsBatch() } +const refreshUpstreamBillingSortedList = async (force = false) => { + if (sortState.sort_by !== 'upstream_billing_rate') return + + const minute = Math.floor(upstreamBillingNow.value / 60_000) + if (!force && lastUpstreamBillingSortRefreshMinute === minute) return + lastUpstreamBillingSortRefreshMinute = minute + try { + await reload() + } catch (error) { + console.error('Failed to refresh upstream billing sort:', error) + } +} + const debouncedReload = () => { syncAccountListDerivedParams() hasPendingListSync.value = false @@ -1033,6 +1017,9 @@ const handleSort = (key: string, order: AccountSortOrder) => { } watch(loading, (isLoading, wasLoading) => { + if (wasLoading && !isLoading) { + upstreamBillingNow.value = Date.now() + } if (wasLoading && !isLoading && pendingTodayStatsRefresh.value) { pendingTodayStatsRefresh.value = false refreshTodayStatsBatch().catch((error) => { @@ -1041,6 +1028,12 @@ watch(loading, (isLoading, wasLoading) => { } }) +watch(upstreamBillingNow, () => { + if (sortState.sort_by !== 'upstream_billing_rate' || loading.value) return + if (typeof document !== 'undefined' && document.hidden) return + void refreshUpstreamBillingSortedList() +}) + const isAnyModalOpen = computed(() => { return ( showCreate.value || @@ -1152,7 +1145,9 @@ const refreshAccountsIncrementally = async () => { pagination.pages = result.data.pages || 0 mergeAccountsIncrementally(result.data.items || []) hasPendingListSync.value = false + markUpstreamBillingSortRefresh() } + upstreamBillingNow.value = Date.now() await refreshTodayStatsBatch() } catch (error) { @@ -1163,11 +1158,20 @@ const refreshAccountsIncrementally = async () => { } const handleManualRefresh = async () => { - await load() + await Promise.all([load(), loadUpstreamBillingProbeGlobalState()]) // Force usage cells to refetch /usage on explicit user refresh. usageManualRefreshToken.value += 1 } +const loadUpstreamBillingProbeGlobalState = async () => { + try { + const settings = await adminAPI.accounts.getUpstreamBillingProbeSettings() + upstreamBillingProbeGloballyEnabled.value = settings.enabled + } catch (error) { + console.error('Failed to load upstream billing probe settings:', error) + } +} + const closeAccountToolsDropdown = () => { showAccountToolsDropdown.value = false } @@ -1197,31 +1201,6 @@ const openTLSFingerprintProfiles = () => { showTLSFingerprintProfiles.value = true } -const loadUpstreamBillingProbeSettings = async () => { - upstreamBillingSettingsLoading.value = true - try { - Object.assign(upstreamBillingProbeSettings, await adminAPI.accounts.getUpstreamBillingProbeSettings()) - } catch (error) { - console.error('Failed to load upstream billing probe settings:', error) - } finally { - upstreamBillingSettingsLoading.value = false - } -} - -const saveUpstreamBillingProbeSettings = async () => { - upstreamBillingSettingsSaving.value = true - try { - const saved = await adminAPI.accounts.updateUpstreamBillingProbeSettings({ ...upstreamBillingProbeSettings }) - Object.assign(upstreamBillingProbeSettings, saved) - appStore.showSuccess(t('admin.accounts.upstreamBilling.settingsSaved')) - } catch (error) { - console.error('Failed to save upstream billing probe settings:', error) - appStore.showError(extractApiErrorMessage(error, t('admin.accounts.upstreamBilling.settingsFailed'))) - } finally { - upstreamBillingSettingsSaving.value = false - } -} - const syncPendingListChanges = async () => { hasPendingListSync.value = false await load() @@ -1397,7 +1376,7 @@ const allColumns = computed(() => { { key: 'priority', label: t('admin.accounts.columns.priority'), sortable: true }, { key: 'scheduler_score', label: t('admin.accounts.columns.schedulerScore'), sortable: false }, { key: 'rate_multiplier', label: t('admin.accounts.columns.billingRateMultiplier'), sortable: true }, - { key: 'upstream_billing_rate', label: t('admin.accounts.columns.upstreamBillingRate'), sortable: false }, + { key: 'upstream_billing_rate', label: t('admin.accounts.columns.upstreamBillingRate'), sortable: true }, { key: 'last_used_at', label: t('admin.accounts.columns.lastUsed'), sortable: true }, { key: 'created_at', label: t('admin.accounts.columns.createdAt'), sortable: true }, { key: 'expires_at', label: t('admin.accounts.columns.expiresAt'), sortable: true }, @@ -1521,9 +1500,14 @@ const handleBulkProbeUpstreamBilling = async () => { accountIDs.forEach(id => probingUpstreamBilling.add(id)) try { const results = await adminAPI.accounts.probeUpstreamBillingBatch(accountIDs) + let patched = false results.forEach(result => { - if (result.snapshot) patchUpstreamBillingSnapshot(result.account_id, result.snapshot) + if (result.snapshot) { + patchUpstreamBillingSnapshot(result.account_id, result.snapshot) + patched = true + } }) + if (patched) await refreshUpstreamBillingSortedList(true) const failed = results.filter(result => result.error).length if (failed > 0) { appStore.showError(t('admin.accounts.upstreamBilling.batchPartial', { success: results.length - failed, failed })) @@ -1789,6 +1773,8 @@ const patchAccountInList = (updatedAccount: Account) => { const patchUpstreamBillingSnapshot = (accountID: number, snapshot: UpstreamBillingProbeSnapshot) => { const account = accounts.value.find(item => item.id === accountID) if (!account) return + markUpstreamBillingSortRefresh() + upstreamBillingNow.value = Date.now() patchAccountInList({ ...account, extra: { ...account.extra, upstream_billing_probe: snapshot } @@ -1799,7 +1785,10 @@ const handleProbeUpstreamBilling = async (account: Account) => { probingUpstreamBilling.add(account.id) try { const result = await adminAPI.accounts.probeUpstreamBilling(account.id) - if (result.snapshot) patchUpstreamBillingSnapshot(account.id, result.snapshot) + if (result.snapshot) { + patchUpstreamBillingSnapshot(account.id, result.snapshot) + await refreshUpstreamBillingSortedList(true) + } } catch (error) { console.error('Failed to probe upstream billing:', error) appStore.showError(extractApiErrorMessage(error, t('admin.accounts.upstreamBilling.probeFailed'))) @@ -2062,7 +2051,7 @@ const handleClickOutside = (event: MouseEvent) => { onMounted(async () => { load() - loadUpstreamBillingProbeSettings() + loadUpstreamBillingProbeGlobalState() try { const [p, g] = await Promise.all([adminAPI.proxies.getAll(), adminAPI.groups.getAll()]) proxies.value = p diff --git a/frontend/src/views/admin/SettingsView.vue b/frontend/src/views/admin/SettingsView.vue index 45ba18336..b2b1c760c 100644 --- a/frontend/src/views/admin/SettingsView.vue +++ b/frontend/src/views/admin/SettingsView.vue @@ -4043,6 +4043,92 @@ </div> </div> + <!-- Upstream Billing Probe Settings --> + <div class="card" data-testid="upstream-billing-probe-settings"> + <div + class="border-b border-gray-100 px-6 py-4 dark:border-dark-700" + > + <h2 class="text-lg font-semibold text-gray-900 dark:text-white"> + {{ t("admin.settings.upstreamBillingProbe.title") }} + </h2> + <p class="mt-1 text-sm text-gray-500 dark:text-gray-400"> + {{ t("admin.settings.upstreamBillingProbe.description") }} + </p> + </div> + <div class="space-y-5 p-6"> + <div + v-if="upstreamBillingProbeLoading" + class="flex items-center gap-2 text-gray-500" + > + <div + class="h-4 w-4 animate-spin rounded-full border-b-2 border-primary-600" + ></div> + {{ t("common.loading") }} + </div> + + <template v-else> + <div class="flex items-center justify-between gap-4"> + <div> + <label class="font-medium text-gray-900 dark:text-white"> + {{ t("admin.settings.upstreamBillingProbe.enabled") }} + </label> + <p class="text-sm text-gray-500 dark:text-gray-400"> + {{ t("admin.settings.upstreamBillingProbe.enabledHint") }} + </p> + </div> + <Toggle + v-model="upstreamBillingProbeForm.enabled" + :aria-label="t('admin.settings.upstreamBillingProbe.enabled')" + data-testid="upstream-billing-probe-enabled" + /> + </div> + + <div + v-if="upstreamBillingProbeForm.enabled" + class="border-t border-gray-100 pt-4 dark:border-dark-700" + > + <label + class="mb-2 block text-sm font-medium text-gray-700 dark:text-gray-300" + for="upstream-billing-probe-interval" + > + {{ t("admin.settings.upstreamBillingProbe.intervalMinutes") }} + </label> + <input + id="upstream-billing-probe-interval" + v-model.number="upstreamBillingProbeForm.interval_minutes" + type="number" + min="5" + max="1440" + class="input w-32" + data-testid="upstream-billing-probe-interval" + @keydown.enter.prevent="saveUpstreamBillingProbeSettings" + /> + <p class="mt-1.5 text-xs text-gray-500 dark:text-gray-400"> + {{ t("admin.settings.upstreamBillingProbe.intervalHint") }} + </p> + </div> + + <div + class="flex justify-end border-t border-gray-100 pt-4 dark:border-dark-700" + > + <button + type="button" + class="btn btn-primary btn-sm" + :disabled="upstreamBillingProbeSaving" + data-testid="upstream-billing-probe-save" + @click="saveUpstreamBillingProbeSettings" + > + {{ + upstreamBillingProbeSaving + ? t("common.saving") + : t("common.save") + }} + </button> + </div> + </template> + </div> + </div> + <!-- Gateway Scheduling Settings --> <div class="card"> <div @@ -7678,6 +7764,14 @@ const adminApiKeyOperating = ref(false); const newAdminApiKey = ref(""); const subscriptionGroups = ref<AdminGroup[]>([]); +// Upstream billing probe state +const upstreamBillingProbeLoading = ref(true); +const upstreamBillingProbeSaving = ref(false); +const upstreamBillingProbeForm = reactive({ + enabled: true, + interval_minutes: 30, +}); + // Overload Cooldown (529) 状态 const overloadCooldownLoading = ref(true); const overloadCooldownSaving = ref(false); @@ -10107,6 +10201,40 @@ function copyNewKey() { }); } +async function loadUpstreamBillingProbeSettings() { + upstreamBillingProbeLoading.value = true; + try { + Object.assign( + upstreamBillingProbeForm, + await adminAPI.accounts.getUpstreamBillingProbeSettings(), + ); + } catch (_error: unknown) { + // Keep defaults when this optional setting cannot be loaded. + } finally { + upstreamBillingProbeLoading.value = false; + } +} + +async function saveUpstreamBillingProbeSettings() { + upstreamBillingProbeSaving.value = true; + try { + const updated = await adminAPI.accounts.updateUpstreamBillingProbeSettings({ + ...upstreamBillingProbeForm, + }); + Object.assign(upstreamBillingProbeForm, updated); + appStore.showSuccess(t("admin.settings.upstreamBillingProbe.saved")); + } catch (error: unknown) { + appStore.showError( + extractApiErrorMessage( + error, + t("admin.settings.upstreamBillingProbe.saveFailed"), + ), + ); + } finally { + upstreamBillingProbeSaving.value = false; + } +} + // Overload Cooldown 方法 async function loadOverloadCooldownSettings() { overloadCooldownLoading.value = true; @@ -10803,6 +10931,7 @@ onMounted(() => { loadSettings(); loadSubscriptionGroups(); loadAdminApiKey(); + loadUpstreamBillingProbeSettings(); loadOverloadCooldownSettings(); loadRateLimit429CooldownSettings(); loadStreamTimeoutSettings(); diff --git a/frontend/src/views/admin/__tests__/AccountsView.bulkEdit.spec.ts b/frontend/src/views/admin/__tests__/AccountsView.bulkEdit.spec.ts index 89ccdf316..39111c53e 100644 --- a/frontend/src/views/admin/__tests__/AccountsView.bulkEdit.spec.ts +++ b/frontend/src/views/admin/__tests__/AccountsView.bulkEdit.spec.ts @@ -7,6 +7,7 @@ const { listAccounts, listWithEtag, getBatchTodayStats, + getUpstreamBillingProbeSettings, getAllProxies, getAllGroups, probeUpstreamBillingBatch @@ -14,6 +15,7 @@ const { listAccounts: vi.fn(), listWithEtag: vi.fn(), getBatchTodayStats: vi.fn(), + getUpstreamBillingProbeSettings: vi.fn(), getAllProxies: vi.fn(), getAllGroups: vi.fn(), probeUpstreamBillingBatch: vi.fn() @@ -25,7 +27,7 @@ vi.mock('@/api/admin', () => ({ list: listAccounts, listWithEtag, getBatchTodayStats, - getUpstreamBillingProbeSettings: vi.fn().mockResolvedValue({ enabled: true, interval_minutes: 30 }), + getUpstreamBillingProbeSettings, delete: vi.fn(), batchClearError: vi.fn(), batchRefresh: vi.fn(), @@ -106,6 +108,7 @@ describe('admin AccountsView bulk edit scope', () => { listAccounts.mockReset() listWithEtag.mockReset() getBatchTodayStats.mockReset() + getUpstreamBillingProbeSettings.mockReset() getAllProxies.mockReset() getAllGroups.mockReset() probeUpstreamBillingBatch.mockReset() @@ -123,6 +126,7 @@ describe('admin AccountsView bulk edit scope', () => { data: null }) getBatchTodayStats.mockResolvedValue({ stats: {} }) + getUpstreamBillingProbeSettings.mockResolvedValue({ enabled: true, interval_minutes: 30 }) getAllProxies.mockResolvedValue([]) getAllGroups.mockResolvedValue([]) probeUpstreamBillingBatch.mockResolvedValue([]) @@ -242,6 +246,75 @@ describe('admin AccountsView bulk edit scope', () => { }) }) + it('passes the loaded global probe state to every upstream billing cell', async () => { + listAccounts.mockResolvedValue({ + items: [ + { + id: 1, + name: 'upstream', + platform: 'openai', + type: 'apikey', + status: 'active', + schedulable: true, + created_at: '2026-07-13T00:00:00Z', + updated_at: '2026-07-13T00:00:00Z' + } + ], + total: 1, + page: 1, + page_size: 20, + pages: 1 + }) + getUpstreamBillingProbeSettings.mockResolvedValue({ enabled: false, interval_minutes: 30 }) + + const wrapper = mount(AccountsView, { + global: { + stubs: { + AppLayout: { template: '<div><slot /></div>' }, + TablePageLayout: { template: '<div><slot name="table" /></div>' }, + DataTable: { + props: ['data'], + template: '<div><div v-for="row in data" :key="row.id"><slot name="cell-upstream_billing_rate" :row="row" /></div></div>' + }, + UpstreamBillingRateCell: { + props: ['globalProbeEnabled'], + template: '<span data-test="upstream-billing-cell" :data-global-enabled="String(globalProbeEnabled)"></span>' + }, + Pagination: true, + ConfirmDialog: true, + AccountTableActions: true, + AccountTableFilters: true, + AccountBulkActionsBar: true, + AccountActionMenu: true, + ImportDataModal: true, + ReAuthAccountModal: true, + AccountTestModal: true, + AccountStatsModal: true, + ScheduledTestsPanel: true, + SyncFromCrsModal: true, + TempUnschedStatusModal: true, + ErrorPassthroughRulesModal: true, + TLSFingerprintProfilesModal: true, + CreateAccountModal: true, + EditAccountModal: true, + BulkEditAccountModal: true, + PlatformTypeBadge: true, + AccountCapacityCell: true, + AccountStatusIndicator: true, + AccountTodayStatsCell: true, + AccountGroupsCell: true, + AccountUsageCell: true, + Icon: true + } + } + }) + + await flushPromises() + + expect(getUpstreamBillingProbeSettings).toHaveBeenCalledTimes(1) + expect(wrapper.get('[data-test="upstream-billing-cell"]').attributes('data-global-enabled')).toBe('false') + }) + it('submits selected account IDs from every page for backend eligibility checks', async () => { const account = (id: number) => ({ id, @@ -302,4 +375,75 @@ describe('admin AccountsView bulk edit scope', () => { expect(probeUpstreamBillingBatch).toHaveBeenCalledWith([7, 11]) }) + + it('reloads the server-sorted list after a batch probe changes a snapshot', async () => { + localStorage.setItem('account-table-sort', JSON.stringify({ key: 'upstream_billing_rate', order: 'asc' })) + const account = (id: number) => ({ + id, + name: `account-${id}`, + platform: 'openai', + type: 'apikey', + status: 'active', + schedulable: true, + created_at: '2026-07-13T00:00:00Z', + updated_at: '2026-07-13T00:00:00Z' + }) + listAccounts + .mockResolvedValueOnce({ items: [account(7)], total: 1, page: 1, page_size: 20, pages: 1 }) + .mockResolvedValueOnce({ items: [account(7)], total: 1, page: 1, page_size: 20, pages: 1 }) + probeUpstreamBillingBatch.mockResolvedValue([ + { + account_id: 7, + snapshot: { + status: 'ok', + data: { effective_rate_multiplier: 0.5 }, + last_attempt_at: '2026-07-13T00:00:00Z', + next_probe_at: '2026-07-13T00:30:00Z' + } + } + ]) + + const wrapper = mount(AccountsView, { + global: { + stubs: { + AppLayout: { template: '<div><slot /></div>' }, + TablePageLayout: { template: '<div><slot name="table" /></div>' }, + DataTable: DataTableStub, + AccountBulkActionsBar: AccountBulkActionsBarStub, + AccountTableActions: true, + AccountTableFilters: true, + AccountActionMenu: true, + Pagination: true, + ConfirmDialog: true, + ImportDataModal: true, + ReAuthAccountModal: true, + AccountTestModal: true, + AccountStatsModal: true, + ScheduledTestsPanel: true, + SyncFromCrsModal: true, + TempUnschedStatusModal: true, + ErrorPassthroughRulesModal: true, + TLSFingerprintProfilesModal: true, + CreateAccountModal: true, + EditAccountModal: true, + BulkEditAccountModal: BulkEditAccountModalStub, + PlatformTypeBadge: true, + AccountCapacityCell: true, + AccountStatusIndicator: true, + AccountTodayStatsCell: true, + AccountGroupsCell: true, + AccountUsageCell: true, + Icon: true + } + } + }) + + await flushPromises() + await wrapper.get('[data-test="select-row"] input').trigger('change') + await wrapper.get('[data-test="probe-upstream-billing"]').trigger('click') + await flushPromises() + + expect(probeUpstreamBillingBatch).toHaveBeenCalledWith([7]) + expect(listAccounts).toHaveBeenCalledTimes(2) + }) }) diff --git a/frontend/src/views/admin/__tests__/AccountsView.usageWindowsHint.spec.ts b/frontend/src/views/admin/__tests__/AccountsView.usageWindowsHint.spec.ts index 0edd426ac..3e8f36f35 100644 --- a/frontend/src/views/admin/__tests__/AccountsView.usageWindowsHint.spec.ts +++ b/frontend/src/views/admin/__tests__/AccountsView.usageWindowsHint.spec.ts @@ -176,5 +176,7 @@ describe('admin AccountsView usage windows hint', () => { expect(wrapper.findAll('[data-test="usage-windows-hint"]').some(node => node.text() === 'admin.accounts.upstreamBilling.trustWarning' )).toBe(true) + const columns = wrapper.getComponent(DataTableStub).props('columns') as Array<{ key: string; sortable: boolean }> + expect(columns.find(column => column.key === 'upstream_billing_rate')?.sortable).toBe(true) }) }) diff --git a/frontend/src/views/admin/__tests__/SettingsView.spec.ts b/frontend/src/views/admin/__tests__/SettingsView.spec.ts index de4666d97..8f7396336 100644 --- a/frontend/src/views/admin/__tests__/SettingsView.spec.ts +++ b/frontend/src/views/admin/__tests__/SettingsView.spec.ts @@ -16,6 +16,8 @@ const { getStreamTimeoutSettings, getRectifierSettings, getBetaPolicySettings, + getUpstreamBillingProbeSettings, + updateUpstreamBillingProbeSettings, getGroups, listProxies, getProviders, @@ -38,6 +40,11 @@ const { getStreamTimeoutSettings: vi.fn(), getRectifierSettings: vi.fn(), getBetaPolicySettings: vi.fn(), + getUpstreamBillingProbeSettings: vi.fn().mockResolvedValue({ + enabled: true, + interval_minutes: 30, + }), + updateUpstreamBillingProbeSettings: vi.fn().mockImplementation(async (payload) => payload), getGroups: vi.fn(), listProxies: vi.fn(), getProviders: vi.fn(), @@ -67,6 +74,10 @@ vi.mock("@/api", () => ({ getRectifierSettings, getBetaPolicySettings, }, + accounts: { + getUpstreamBillingProbeSettings, + updateUpstreamBillingProbeSettings, + }, groups: { getAll: getGroups, }, @@ -184,6 +195,14 @@ vi.mock("vue-i18n", async () => { "admin.settings.openaiExperimentalScheduler.upstreamCostWeight": "计费倍率", "admin.settings.openaiExperimentalScheduler.previousResponseWeight": "previous_response 粘性", "admin.settings.openaiExperimentalScheduler.sessionStickyWeight": "session_hash 粘性", + "admin.settings.upstreamBillingProbe.title": "上游倍率自动探测", + "admin.settings.upstreamBillingProbe.description": "定期获取 OpenAI API Key 所连接上游 Sub2API 站点声明的计费倍率。", + "admin.settings.upstreamBillingProbe.enabled": "启用全局自动探测", + "admin.settings.upstreamBillingProbe.enabledHint": "开启后,仅对账号自身已启用自动检测的账号执行定时探测。", + "admin.settings.upstreamBillingProbe.intervalMinutes": "探测周期(分钟)", + "admin.settings.upstreamBillingProbe.intervalHint": "范围 5–1440 分钟。", + "admin.settings.upstreamBillingProbe.saved": "上游倍率自动探测设置已保存", + "admin.settings.upstreamBillingProbe.saveFailed": "保存上游倍率自动探测设置失败", "admin.settings.site.uploadImage": "上传图片", "admin.settings.site.remove": "移除", "admin.settings.platformQuota.platform": "平台", @@ -517,6 +536,16 @@ async function openSecurityTab(wrapper: ReturnType<typeof mountView>) { await flushPromises(); } +async function openGatewayTab(wrapper: ReturnType<typeof mountView>) { + const gatewayTabButton = wrapper + .findAll("button") + .find((node) => node.text().includes("admin.settings.tabs.gateway")); + + expect(gatewayTabButton).toBeDefined(); + await gatewayTabButton?.trigger("click"); + await flushPromises(); +} + async function openUsersTab(wrapper: ReturnType<typeof mountView>) { const usersTabButton = wrapper .findAll("button") @@ -540,6 +569,8 @@ describe("admin SettingsView payment visible method controls", () => { getStreamTimeoutSettings.mockReset(); getRectifierSettings.mockReset(); getBetaPolicySettings.mockReset(); + getUpstreamBillingProbeSettings.mockReset(); + updateUpstreamBillingProbeSettings.mockReset(); getGroups.mockReset(); listProxies.mockReset(); getProviders.mockReset(); @@ -595,6 +626,11 @@ describe("admin SettingsView payment visible method controls", () => { getBetaPolicySettings.mockResolvedValue({ rules: [], }); + getUpstreamBillingProbeSettings.mockResolvedValue({ + enabled: true, + interval_minutes: 30, + }); + updateUpstreamBillingProbeSettings.mockImplementation(async (payload) => payload); getGroups.mockResolvedValue([]); listProxies.mockResolvedValue({ items: [], @@ -848,6 +884,38 @@ describe("admin SettingsView payment visible method controls", () => { expect(wrapper.text()).not.toContain("OpenAI 高级调度器"); }); + it("loads and saves upstream billing probe settings from the gateway tab", async () => { + getUpstreamBillingProbeSettings.mockResolvedValueOnce({ + enabled: false, + interval_minutes: 45, + }); + + const wrapper = mountView(); + + await flushPromises(); + await openGatewayTab(wrapper); + + const card = wrapper.get('[data-testid="upstream-billing-probe-settings"]'); + expect(card.isVisible()).toBe(true); + expect(card.text()).toContain("上游倍率自动探测"); + expect( + (card.get('[data-testid="upstream-billing-probe-enabled"]').element as HTMLInputElement) + .checked, + ).toBe(false); + expect(card.find('[data-testid="upstream-billing-probe-interval"]').exists()).toBe(false); + + await card.get('[data-testid="upstream-billing-probe-enabled"]').setValue(true); + await card.get('[data-testid="upstream-billing-probe-interval"]').setValue(60); + await card.get('[data-testid="upstream-billing-probe-save"]').trigger("click"); + await flushPromises(); + + expect(updateUpstreamBillingProbeSettings).toHaveBeenCalledWith({ + enabled: true, + interval_minutes: 60, + }); + expect(showSuccess).toHaveBeenCalledWith("上游倍率自动探测设置已保存"); + }); + it("places and explains rate controls for both scheduling modes", async () => { const wrapper = mountView(); diff --git a/frontend/src/views/admin/ops/components/OpsErrorDetailModal.vue b/frontend/src/views/admin/ops/components/OpsErrorDetailModal.vue index c346c547b..348eba297 100644 --- a/frontend/src/views/admin/ops/components/OpsErrorDetailModal.vue +++ b/frontend/src/views/admin/ops/components/OpsErrorDetailModal.vue @@ -114,23 +114,6 @@ </div> </div> - <div v-if="detail.attempted_key_prefix" class="rounded-xl bg-gray-50 p-4 dark:bg-dark-900"> - <div class="text-xs font-bold uppercase tracking-wider text-gray-400">{{ t('admin.ops.errorDetail.attemptedKeyPrefix') }}</div> - <div class="mt-1 font-mono text-sm font-medium text-gray-900 dark:text-white"> - {{ detail.attempted_key_prefix }} - </div> - </div> - - <div v-if="detail.deleted_key_owner_email" class="rounded-xl bg-gray-50 p-4 dark:bg-dark-900"> - <div class="text-xs font-bold uppercase tracking-wider text-gray-400">{{ t('admin.ops.errorDetail.deletedKeyOwner') }}</div> - <div class="mt-1 text-sm font-medium text-gray-900 dark:text-white"> - {{ detail.deleted_key_owner_email }} - <span v-if="detail.deleted_key_name" class="ml-1 text-xs text-gray-500 dark:text-gray-400">({{ detail.deleted_key_name }})</span> - <span class="ml-2 inline-flex items-center rounded px-1.5 py-0.5 text-[10px] font-bold ring-1 ring-inset bg-red-50 text-red-700 ring-red-600/20 dark:bg-red-900/30 dark:text-red-400 dark:ring-red-500/30"> - {{ t('admin.ops.errorDetail.keyDeletedBadge') }} - </span> - </div> - </div> </div> <!-- Response content (client request -> error_body; upstream -> upstream_error_detail/message) --> diff --git a/frontend/src/views/admin/ops/components/OpsErrorLogTable.vue b/frontend/src/views/admin/ops/components/OpsErrorLogTable.vue index 6cb0bb1b4..e8efc2e87 100644 --- a/frontend/src/views/admin/ops/components/OpsErrorLogTable.vue +++ b/frontend/src/views/admin/ops/components/OpsErrorLogTable.vue @@ -77,19 +77,6 @@ <span v-else class="font-medium text-gray-900 dark:text-white">{{ row.user_email || '-' }}</span> <span class="ml-1 text-gray-500 dark:text-gray-400">#{{ row.user_id }}</span> </div> - <!-- 认证失败行 user_id 为空:回退显示已删除 KEY 所有者(归因快照,与详情弹窗一致) --> - <div v-else-if="row.deleted_key_owner_user_id" class="text-sm"> - <button - v-if="userClickable && row.deleted_key_owner_email" - class="font-medium text-primary-600 underline decoration-dashed underline-offset-2 transition-colors hover:text-primary-700 dark:text-primary-400 dark:hover:text-primary-300" - :title="t('admin.usage.clickToViewBalance')" - @click.stop="emit('userClick', row.deleted_key_owner_user_id, row.deleted_key_owner_email ?? undefined)" - > - {{ row.deleted_key_owner_email }} - </button> - <span v-else class="font-medium text-gray-900 dark:text-white">{{ row.deleted_key_owner_email || '-' }}</span> - <span class="ml-1 text-gray-500 dark:text-gray-400">#{{ row.deleted_key_owner_user_id }}</span> - </div> <span v-else class="text-sm text-gray-400 dark:text-gray-500">-</span> </template> diff --git a/frontend/src/views/admin/ops/components/OpsSettingsDialog.vue b/frontend/src/views/admin/ops/components/OpsSettingsDialog.vue index bfb7a65f8..105ea560a 100644 --- a/frontend/src/views/admin/ops/components/OpsSettingsDialog.vue +++ b/frontend/src/views/admin/ops/components/OpsSettingsDialog.vue @@ -572,16 +572,6 @@ async function saveAllSettings() { <Toggle v-model="advancedSettings.ignore_no_available_accounts" /> </div> - <div class="flex items-center justify-between"> - <div> - <label class="text-sm font-medium text-gray-700 dark:text-gray-300">{{ t('admin.ops.settings.ignoreInvalidApiKeyErrors') }}</label> - <p class="mt-1 text-xs text-gray-500"> - {{ t('admin.ops.settings.ignoreInvalidApiKeyErrorsHint') }} - </p> - </div> - <Toggle v-model="advancedSettings.ignore_invalid_api_key_errors" /> - </div> - <div class="flex items-center justify-between"> <div> <label class="text-sm font-medium text-gray-700 dark:text-gray-300">{{ t('admin.ops.settings.ignoreInsufficientBalanceErrors') }}</label> diff --git a/frontend/src/views/admin/orders/AdminPaymentPlansView.vue b/frontend/src/views/admin/orders/AdminPaymentPlansView.vue index eb4f39a31..17a17a5f1 100644 --- a/frontend/src/views/admin/orders/AdminPaymentPlansView.vue +++ b/frontend/src/views/admin/orders/AdminPaymentPlansView.vue @@ -142,7 +142,7 @@ const planColumns = computed((): Column[] => [ { key: 'name', label: t('payment.admin.planName') }, { key: 'group_id', label: t('payment.admin.group') }, { key: 'price', label: t('payment.admin.price') }, - { key: 'validity_days', label: t('payment.admin.validityDays') }, + { key: 'validity_days', label: t('payment.admin.validity') }, { key: 'for_sale', label: t('payment.admin.forSale') }, { key: 'sort_order', label: t('payment.admin.sortOrder') }, { key: 'actions', label: t('common.actions') }, diff --git a/frontend/src/views/admin/orders/PlanEditDialog.vue b/frontend/src/views/admin/orders/PlanEditDialog.vue index 0ddc4cab6..7535ebc56 100644 --- a/frontend/src/views/admin/orders/PlanEditDialog.vue +++ b/frontend/src/views/admin/orders/PlanEditDialog.vue @@ -48,7 +48,7 @@ <div><label class="input-label">{{ t('payment.admin.originalPrice') }}</label><input v-model.number="planForm.original_price" type="number" step="0.01" min="0" class="input" /></div> </div> <div class="grid grid-cols-2 gap-4"> - <div><label class="input-label">{{ t('payment.admin.validityDays') }} <span class="text-red-500">*</span></label><input v-model.number="planForm.validity_days" type="number" min="1" class="input" required /></div> + <div><label class="input-label">{{ t('payment.admin.validity') }} <span class="text-red-500">*</span></label><input v-model.number="planForm.validity_days" type="number" min="1" class="input" required /></div> <div><label class="input-label">{{ t('payment.admin.validityUnit') }} <span class="text-red-500">*</span></label><Select v-model="planForm.validity_unit" :options="validityUnitOptions" /></div> </div> <div class="grid grid-cols-2 gap-4"> @@ -211,7 +211,7 @@ async function handleSavePlan() { return } if (!planForm.validity_days || planForm.validity_days < 1) { - appStore.showError(t('payment.admin.validityDaysRequired')) + appStore.showError(t('payment.admin.validityRequired')) return } saving.value = true diff --git a/frontend/src/views/public/LegalDocumentView.vue b/frontend/src/views/public/LegalDocumentView.vue index 601bb12d2..298f23a0b 100644 --- a/frontend/src/views/public/LegalDocumentView.vue +++ b/frontend/src/views/public/LegalDocumentView.vue @@ -3,12 +3,18 @@ <header class="border-b border-gray-200 bg-white/95 dark:border-dark-800 dark:bg-dark-900/95"> <div class="mx-auto flex max-w-5xl items-center justify-between gap-4 px-4 py-4 sm:px-6"> <RouterLink to="/home" class="flex min-w-0 items-center gap-3"> - <span class="flex h-10 w-10 flex-shrink-0 items-center justify-center overflow-hidden rounded-xl bg-white shadow-sm ring-1 ring-gray-200 dark:bg-dark-800 dark:ring-dark-700"> - <img :src="siteLogo || '/logo.png'" alt="Logo" class="h-full w-full object-contain" /> - </span> - <span class="truncate text-base font-semibold text-gray-950 dark:text-white"> - {{ siteName }} - </span> + <template v-if="settings"> + <span class="flex h-10 w-10 flex-shrink-0 items-center justify-center overflow-hidden rounded-xl bg-white shadow-sm ring-1 ring-gray-200 dark:bg-dark-800 dark:ring-dark-700"> + <img :src="siteLogo || '/logo.png'" alt="Logo" class="h-full w-full object-contain" /> + </span> + <span class="truncate text-base font-semibold text-gray-950 dark:text-white"> + {{ siteName }} + </span> + </template> + <template v-else> + <span class="h-10 w-10 flex-shrink-0 animate-pulse rounded-xl bg-gray-200 dark:bg-dark-700" aria-hidden="true"></span> + <span class="h-5 w-28 animate-pulse rounded bg-gray-200 dark:bg-dark-700" aria-hidden="true"></span> + </template> </RouterLink> <RouterLink to="/login" @@ -90,10 +96,10 @@ import { marked } from 'marked' import DOMPurify from 'dompurify' import { useI18n } from 'vue-i18n' import Icon from '@/components/icons/Icon.vue' -import { getPublicSettings } from '@/api/auth' import { getLocale } from '@/i18n' import { sanitizeUrl } from '@/utils/url' -import type { LoginAgreementDocument, PublicSettings } from '@/types' +import { useAppStore } from '@/stores/app' +import type { LoginAgreementDocument } from '@/types' import zhAdminCompliance from '../../../../docs/legal/admin-compliance.zh.md?raw' import enAdminCompliance from '../../../../docs/legal/admin-compliance.en.md?raw' @@ -101,8 +107,9 @@ type LegalDocumentIcon = 'document' | 'shield' | 'globe' | 'cog' const route = useRoute() const { t } = useI18n() -const settings = ref<PublicSettings | null>(null) -const loading = ref(true) +const appStore = useAppStore() +const settings = computed(() => appStore.cachedPublicSettings) +const loading = ref(!settings.value) const loadError = ref(false) marked.setOptions({ @@ -166,15 +173,12 @@ const documentIcon = computed<LegalDocumentIcon>(() => { }) onMounted(async () => { - loading.value = true loadError.value = false - try { - settings.value = await getPublicSettings() - } catch { + const loadedSettings = await appStore.fetchPublicSettings() + if (!loadedSettings) { loadError.value = true - } finally { - loading.value = false } + loading.value = false }) </script> diff --git a/frontend/vite.config.ts b/frontend/vite.config.ts index 303a4bf77..5df99b8fa 100644 --- a/frontend/vite.config.ts +++ b/frontend/vite.config.ts @@ -3,6 +3,49 @@ import vue from '@vitejs/plugin-vue' import checker from 'vite-plugin-checker' import { resolve } from 'path' +function escapeHtml(value: string): string { + return value.replace(/[&<>"']/g, (character) => ({ + '&': '&', + '<': '<', + '>': '>', + '"': '"', + "'": ''', + })[character] || character) +} + +function isSafeImageUrl(value: string): boolean { + const trimmed = value.trim() + if ((trimmed.startsWith('/') && !trimmed.startsWith('//')) || /^data:image\//i.test(trimmed)) { + return true + } + try { + const parsed = new URL(trimmed) + return parsed.protocol === 'http:' || parsed.protocol === 'https:' + } catch { + return false + } +} + +function injectBranding(html: string, config: { site_name?: string; site_logo?: string }): string { + let brandedHtml = html + const siteName = config.site_name?.trim() + if (siteName) { + brandedHtml = brandedHtml.replace( + /<title>[^<]*<\/title>/i, + `<title>${escapeHtml(siteName)} - AI API Gateway`, + ) + } + + const siteLogo = config.site_logo?.trim() + if (siteLogo && isSafeImageUrl(siteLogo)) { + brandedHtml = brandedHtml.replace( + /]*>/i, + ``, + ) + } + return brandedHtml +} + /** * Vite 插件:开发模式下注入公开配置到 index.html * 与生产模式的后端注入行为保持一致,消除闪烁 @@ -22,7 +65,7 @@ function injectPublicSettings(backendUrl: string): Plugin { const data = await response.json() if (data.code === 0 && data.data) { const script = `` - return html.replace('', `${script}\n`) + return injectBranding(html, data.data).replace('', `${script}\n`) } } } catch (e) {