diff --git a/deploy/config.yaml b/deploy/config.yaml index 98d8396..e3f36f5 100644 --- a/deploy/config.yaml +++ b/deploy/config.yaml @@ -78,6 +78,7 @@ rabbitmq: compliance_review_dlq: compliance.review.run.dlq compliance_review_dlq_routing_key: compliance.review.run.dlq consumer_prefetch: 10 + generation_consumer_prefetch: 1 publish_channel_pool_size: 32 scheduler: @@ -218,6 +219,16 @@ llm: temperature: 0.7 reasoning_effort: minimal web_search_limit: 3 + # Global Redis-backed budget shared by all replicas. + rate_limit_scope: global + rate_limit_key: ark:doubao-seed-2-0 + rate_limit_failover_mode: process_degraded + rate_limit_failover_replicas: 10 + request_rate_limit_rpm: 30000 + request_rate_limit_burst: 3000 + token_rate_limit_tpm: 5000000 + token_rate_limit_burst: 500000 + estimated_chars_per_token: 1 compliance: enabled: true @@ -246,7 +257,7 @@ retrieval: generation: queue_size: 128 - worker_concurrency: 8 + worker_concurrency: 24 stream_enabled: false article_timeout: 8m task_lease_ttl: 10m diff --git a/deploy/k3s/apps.yaml b/deploy/k3s/apps.yaml index 043d8ee..b7e5e33 100644 --- a/deploy/k3s/apps.yaml +++ b/deploy/k3s/apps.yaml @@ -206,7 +206,7 @@ metadata: labels: app.kubernetes.io/name: worker-generate spec: - replicas: 1 + replicas: 3 selector: matchLabels: app.kubernetes.io/name: worker-generate diff --git a/deploy/k3s/config/config.yaml b/deploy/k3s/config/config.yaml index 8bff38c..7eb6e6f 100644 --- a/deploy/k3s/config/config.yaml +++ b/deploy/k3s/config/config.yaml @@ -83,6 +83,7 @@ rabbitmq: compliance_review_dlq: compliance.review.run.dlq compliance_review_dlq_routing_key: compliance.review.run.dlq consumer_prefetch: 10 + generation_consumer_prefetch: 1 publish_channel_pool_size: 32 scheduler: @@ -200,6 +201,16 @@ llm: temperature: 0.7 reasoning_effort: minimal web_search_limit: 3 + # Global Redis-backed budget shared by all replicas. + rate_limit_scope: global + rate_limit_key: ark:doubao-seed-2-0 + rate_limit_failover_mode: process_degraded + rate_limit_failover_replicas: 10 + request_rate_limit_rpm: 30000 + request_rate_limit_burst: 3000 + token_rate_limit_tpm: 5000000 + token_rate_limit_burst: 500000 + estimated_chars_per_token: 1 compliance: enabled: true @@ -228,7 +239,7 @@ retrieval: generation: queue_size: 128 - worker_concurrency: 8 + worker_concurrency: 24 stream_enabled: false article_timeout: 8m task_lease_ttl: 10m diff --git a/server/configs/config.local.yaml.example b/server/configs/config.local.yaml.example index 0815e07..f81d618 100644 --- a/server/configs/config.local.yaml.example +++ b/server/configs/config.local.yaml.example @@ -70,6 +70,7 @@ rabbitmq: compliance_review_dlq: compliance.review.run.dlq compliance_review_dlq_routing_key: compliance.review.run.dlq consumer_prefetch: 10 + generation_consumer_prefetch: 1 publish_channel_pool_size: 32 # Prefer environment override when needed: # export RABBITMQ_URL=amqp://geo:geo_dev@localhost:5672/geo @@ -198,6 +199,16 @@ llm: temperature: 0.7 reasoning_effort: minimal web_search_limit: 3 + # Global Redis-backed budget shared by all replicas. + rate_limit_scope: global + rate_limit_key: ark:doubao-seed-2-0 + rate_limit_failover_mode: process_degraded + rate_limit_failover_replicas: 10 + request_rate_limit_rpm: 30000 + request_rate_limit_burst: 3000 + token_rate_limit_tpm: 5000000 + token_rate_limit_burst: 500000 + estimated_chars_per_token: 1 # top_p: 0.95 # Prefer environment override for secrets: # export LLM_API_KEY=your-api-key @@ -233,7 +244,7 @@ retrieval: generation: queue_size: 128 - worker_concurrency: 8 + worker_concurrency: 24 stream_enabled: false article_timeout: 8m diff --git a/server/configs/config.yaml b/server/configs/config.yaml index 0abfd4e..f46a082 100644 --- a/server/configs/config.yaml +++ b/server/configs/config.yaml @@ -83,6 +83,7 @@ rabbitmq: compliance_review_dlq: compliance.review.run.dlq compliance_review_dlq_routing_key: compliance.review.run.dlq consumer_prefetch: 10 + generation_consumer_prefetch: 1 publish_channel_pool_size: 32 scheduler: @@ -204,6 +205,16 @@ llm: temperature: 0.7 reasoning_effort: minimal web_search_limit: 3 + # Global Redis-backed budget shared by all replicas. + rate_limit_scope: global + rate_limit_key: ark:doubao-seed-2-0 + rate_limit_failover_mode: process_degraded + rate_limit_failover_replicas: 10 + request_rate_limit_rpm: 30000 + request_rate_limit_burst: 3000 + token_rate_limit_tpm: 5000000 + token_rate_limit_burst: 500000 + estimated_chars_per_token: 1 compliance: enabled: true @@ -232,7 +243,7 @@ retrieval: generation: queue_size: 128 - worker_concurrency: 8 + worker_concurrency: 24 stream_enabled: false article_timeout: 8m task_lease_ttl: 10m diff --git a/server/go.mod b/server/go.mod index ecbfe66..4fb1deb 100644 --- a/server/go.mod +++ b/server/go.mod @@ -33,6 +33,7 @@ require ( golang.org/x/image v0.38.0 golang.org/x/net v0.48.0 golang.org/x/sync v0.20.0 + golang.org/x/time v0.5.0 gopkg.in/yaml.v3 v3.0.1 ) @@ -83,7 +84,6 @@ require ( golang.org/x/arch v0.3.0 // indirect golang.org/x/sys v0.39.0 // indirect golang.org/x/text v0.35.0 // indirect - golang.org/x/time v0.5.0 // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20251202230838-ff82c1b0f217 // indirect google.golang.org/grpc v1.79.3 // indirect google.golang.org/protobuf v1.36.10 // indirect diff --git a/server/internal/bootstrap/bootstrap.go b/server/internal/bootstrap/bootstrap.go index def4825..04d4655 100644 --- a/server/internal/bootstrap/bootstrap.go +++ b/server/internal/bootstrap/bootstrap.go @@ -135,7 +135,7 @@ func New(configPath string) (*App, error) { logger.Warn("tenant login guard disabled by config; brute-force / rate-limit protection is off") } loginGuard := auth.NewLoginGuard(rdb, loginGuardCfg) - llmClient := llm.NewReloadableClient(llm.New(cfg.LLM)) + llmClient := llm.NewReloadableClient(llm.NewWithRedis(cfg.LLM, rdb)) retrievalProvider := retrieval.NewReloadableProvider(retrieval.NewProvider(cfg.Retrieval, logger)) vectorStore := retrieval.NewReloadableVectorStore(retrieval.NewQdrantStore(cfg.Qdrant, logger)) objectStorageClient := objectstorage.NewReloadableClient(objectstorage.New(cfg.ObjectStorage, logger)) @@ -294,7 +294,7 @@ func (a *App) applyReloadedConfig(event config.ReloadEvent) { a.JWT.Update(event.Current.JWT.Secret, event.Current.JWT.AccessTTL, event.Current.JWT.RefreshTTL) if a.ReloadableLLM != nil { - a.ReloadableLLM.Update(llm.New(event.Current.LLM)) + a.ReloadableLLM.Update(llm.NewWithRedis(event.Current.LLM, a.Redis)) } if a.ReloadableRetrieval != nil { a.ReloadableRetrieval.Update(retrieval.NewProvider(event.Current.Retrieval, a.Logger)) diff --git a/server/internal/shared/bootstrap/http_test.go b/server/internal/shared/bootstrap/http_test.go index 08587d1..4e89084 100644 --- a/server/internal/shared/bootstrap/http_test.go +++ b/server/internal/shared/bootstrap/http_test.go @@ -12,8 +12,6 @@ import ( ) func TestNewEngineUsesForwardedClientIPFromTrustedProxy(t *testing.T) { - t.Parallel() - gin.SetMode(gin.TestMode) engine, err := NewEngine(config.ServerConfig{ @@ -38,8 +36,6 @@ func TestNewEngineUsesForwardedClientIPFromTrustedProxy(t *testing.T) { } func TestNewEngineRejectsInvalidTrustedProxy(t *testing.T) { - t.Parallel() - gin.SetMode(gin.TestMode) _, err := NewEngine(config.ServerConfig{ diff --git a/server/internal/shared/config/config.go b/server/internal/shared/config/config.go index 2a25cdb..8e0e3de 100644 --- a/server/internal/shared/config/config.go +++ b/server/internal/shared/config/config.go @@ -158,6 +158,7 @@ type RabbitMQConfig struct { ComplianceReviewDLQ string `mapstructure:"compliance_review_dlq"` ComplianceReviewDLQRouteKey string `mapstructure:"compliance_review_dlq_routing_key"` ConsumerPrefetch int `mapstructure:"consumer_prefetch"` + GenerationConsumerPrefetch int `mapstructure:"generation_consumer_prefetch"` PublishChannelPoolSize int `mapstructure:"publish_channel_pool_size"` } @@ -340,6 +341,15 @@ type LLMConfig struct { TopP float64 `mapstructure:"top_p"` ReasoningEffort string `mapstructure:"reasoning_effort"` WebSearchLimit int32 `mapstructure:"web_search_limit"` + RequestRateLimitRPM int `mapstructure:"request_rate_limit_rpm"` + RequestRateLimitBurst int `mapstructure:"request_rate_limit_burst"` + TokenRateLimitTPM int64 `mapstructure:"token_rate_limit_tpm"` + TokenRateLimitBurst int `mapstructure:"token_rate_limit_burst"` + EstimatedCharsPerToken float64 `mapstructure:"estimated_chars_per_token"` + RateLimitScope string `mapstructure:"rate_limit_scope"` + RateLimitKey string `mapstructure:"rate_limit_key"` + RateLimitFailoverMode string `mapstructure:"rate_limit_failover_mode"` + RateLimitFailoverReplicas int `mapstructure:"rate_limit_failover_replicas"` } type ComplianceConfig struct { @@ -480,6 +490,7 @@ func loadWithFiles(configPath string) (*Config, []string, error) { normalizeMonitoringDispatchConfig(&cfg.MonitoringDispatch) normalizeMembershipConfig(&cfg.Membership) normalizeBrandLibraryConfig(&cfg.BrandLibrary) + NormalizeLLMConfig(&cfg.LLM) NormalizeComplianceConfig(&cfg.Compliance) NormalizeGenerationConfig(&cfg.Generation) NormalizeBrowserFetchConfig(&cfg.BrowserFetch) @@ -924,6 +935,47 @@ func NormalizeComplianceConfig(cfg *ComplianceConfig) { } } +func NormalizeLLMConfig(cfg *LLMConfig) { + if cfg == nil { + return + } + if cfg.RequestRateLimitRPM < 0 { + cfg.RequestRateLimitRPM = 0 + } + if cfg.RequestRateLimitBurst < 0 { + cfg.RequestRateLimitBurst = 0 + } + if cfg.TokenRateLimitTPM < 0 { + cfg.TokenRateLimitTPM = 0 + } + if cfg.TokenRateLimitBurst < 0 { + cfg.TokenRateLimitBurst = 0 + } + if cfg.EstimatedCharsPerToken <= 0 { + cfg.EstimatedCharsPerToken = 1 + } + switch strings.ToLower(strings.TrimSpace(cfg.RateLimitScope)) { + case "", "process", "local": + cfg.RateLimitScope = "process" + case "global", "redis", "distributed": + cfg.RateLimitScope = "global" + default: + cfg.RateLimitScope = "process" + } + cfg.RateLimitKey = strings.TrimSpace(cfg.RateLimitKey) + switch strings.ToLower(strings.TrimSpace(cfg.RateLimitFailoverMode)) { + case "", "fail_closed", "closed": + cfg.RateLimitFailoverMode = "fail_closed" + case "process_degraded", "local_degraded", "process", "local": + cfg.RateLimitFailoverMode = "process_degraded" + default: + cfg.RateLimitFailoverMode = "fail_closed" + } + if cfg.RateLimitFailoverReplicas < 0 { + cfg.RateLimitFailoverReplicas = 0 + } +} + func candidateConfigPaths(configPath string, local bool) []string { trimmed := strings.TrimSpace(configPath) if trimmed == "" { @@ -992,6 +1044,33 @@ func applyEnvOverrides(cfg *Config) { if model, ok := lookupNonEmptyEnv("LLM_MONITORING_ANSWER_PARSE_MODEL"); ok { cfg.LLM.MonitoringAnswerParseModel = model } + if n, ok := lookupIntEnv("LLM_REQUEST_RATE_LIMIT_RPM"); ok { + cfg.LLM.RequestRateLimitRPM = n + } + if n, ok := lookupIntEnv("LLM_REQUEST_RATE_LIMIT_BURST"); ok { + cfg.LLM.RequestRateLimitBurst = n + } + if n, ok := lookupInt64Env("LLM_TOKEN_RATE_LIMIT_TPM"); ok { + cfg.LLM.TokenRateLimitTPM = n + } + if n, ok := lookupIntEnv("LLM_TOKEN_RATE_LIMIT_BURST"); ok { + cfg.LLM.TokenRateLimitBurst = n + } + if n, ok := lookupFloatEnv("LLM_ESTIMATED_CHARS_PER_TOKEN"); ok { + cfg.LLM.EstimatedCharsPerToken = n + } + if scope, ok := lookupNonEmptyEnv("LLM_RATE_LIMIT_SCOPE"); ok { + cfg.LLM.RateLimitScope = scope + } + if key, ok := lookupNonEmptyEnv("LLM_RATE_LIMIT_KEY"); ok { + cfg.LLM.RateLimitKey = key + } + if mode, ok := lookupNonEmptyEnv("LLM_RATE_LIMIT_FAILOVER_MODE"); ok { + cfg.LLM.RateLimitFailoverMode = mode + } + if n, ok := lookupIntEnv("LLM_RATE_LIMIT_FAILOVER_REPLICAS"); ok { + cfg.LLM.RateLimitFailoverReplicas = n + } if enabled, ok := lookupBoolEnv("COMPLIANCE_ENABLED"); ok { cfg.Compliance.Enabled = enabled } @@ -1136,6 +1215,12 @@ func applyEnvOverrides(cfg *Config) { if rabbitmqURL, ok := lookupNonEmptyEnv("RABBITMQ_URL"); ok { cfg.RabbitMQ.URL = rabbitmqURL } + if n, ok := lookupIntEnv("RABBITMQ_CONSUMER_PREFETCH"); ok { + cfg.RabbitMQ.ConsumerPrefetch = n + } + if n, ok := lookupIntEnv("RABBITMQ_GENERATION_CONSUMER_PREFETCH"); ok { + cfg.RabbitMQ.GenerationConsumerPrefetch = n + } if url, ok := lookupNonEmptyEnv("QDRANT_URL"); ok { cfg.Qdrant.URL = url } @@ -1391,6 +1476,12 @@ func normalizeRabbitMQConfig(cfg *RabbitMQConfig) { if strings.TrimSpace(cfg.ComplianceReviewDLQRouteKey) == "" { cfg.ComplianceReviewDLQRouteKey = "compliance.review.run.dlq" } + if cfg.ConsumerPrefetch < 0 { + cfg.ConsumerPrefetch = 0 + } + if cfg.GenerationConsumerPrefetch < 0 { + cfg.GenerationConsumerPrefetch = 0 + } if cfg.PublishChannelPoolSize <= 0 { cfg.PublishChannelPoolSize = 16 } diff --git a/server/internal/shared/config/config_test.go b/server/internal/shared/config/config_test.go index 7f63c6c..f2bc07c 100644 --- a/server/internal/shared/config/config_test.go +++ b/server/internal/shared/config/config_test.go @@ -51,6 +51,125 @@ compliance: } } +func TestLoadAppliesLLMRateLimitConfig(t *testing.T) { + configPath := writeTestConfig(t, ` +llm: + rate_limit_scope: global + rate_limit_key: ark:shared + request_rate_limit_rpm: 30000 + request_rate_limit_burst: 100 + token_rate_limit_tpm: 5000000 + token_rate_limit_burst: 166000 + estimated_chars_per_token: 1.25 + rate_limit_failover_mode: process_degraded + rate_limit_failover_replicas: 10 +`) + + cfg, err := Load(configPath) + if err != nil { + t.Fatalf("Load() error = %v", err) + } + + if cfg.LLM.RateLimitScope != "global" || + cfg.LLM.RateLimitKey != "ark:shared" || + cfg.LLM.RequestRateLimitRPM != 30000 || + cfg.LLM.RequestRateLimitBurst != 100 || + cfg.LLM.TokenRateLimitTPM != 5000000 || + cfg.LLM.TokenRateLimitBurst != 166000 || + cfg.LLM.EstimatedCharsPerToken != 1.25 || + cfg.LLM.RateLimitFailoverMode != "process_degraded" || + cfg.LLM.RateLimitFailoverReplicas != 10 { + t.Fatalf("unexpected llm rate limit config: %#v", cfg.LLM) + } +} + +func TestLoadAppliesLLMRateLimitEnvOverrides(t *testing.T) { + t.Setenv("LLM_RATE_LIMIT_SCOPE", "global") + t.Setenv("LLM_RATE_LIMIT_KEY", "env-ark-shared") + t.Setenv("LLM_REQUEST_RATE_LIMIT_RPM", "9000") + t.Setenv("LLM_REQUEST_RATE_LIMIT_BURST", "90") + t.Setenv("LLM_TOKEN_RATE_LIMIT_TPM", "1500000") + t.Setenv("LLM_TOKEN_RATE_LIMIT_BURST", "150000") + t.Setenv("LLM_ESTIMATED_CHARS_PER_TOKEN", "1.5") + t.Setenv("LLM_RATE_LIMIT_FAILOVER_MODE", "local") + t.Setenv("LLM_RATE_LIMIT_FAILOVER_REPLICAS", "12") + + configPath := writeTestConfig(t, ` +llm: + rate_limit_scope: process + request_rate_limit_rpm: 10000 + token_rate_limit_tpm: 1660000 +`) + + cfg, err := Load(configPath) + if err != nil { + t.Fatalf("Load() error = %v", err) + } + + if cfg.LLM.RateLimitScope != "global" || + cfg.LLM.RateLimitKey != "env-ark-shared" || + cfg.LLM.RequestRateLimitRPM != 9000 || + cfg.LLM.RequestRateLimitBurst != 90 || + cfg.LLM.TokenRateLimitTPM != 1500000 || + cfg.LLM.TokenRateLimitBurst != 150000 || + cfg.LLM.EstimatedCharsPerToken != 1.5 || + cfg.LLM.RateLimitFailoverMode != "process_degraded" || + cfg.LLM.RateLimitFailoverReplicas != 12 { + t.Fatalf("unexpected llm env rate limit config: %#v", cfg.LLM) + } +} + +func TestLoadNormalizesLLMRateLimitConfig(t *testing.T) { + configPath := writeTestConfig(t, ` +llm: + request_rate_limit_rpm: -1 + request_rate_limit_burst: -1 + token_rate_limit_tpm: -1 + token_rate_limit_burst: -1 + estimated_chars_per_token: 0 + rate_limit_scope: broken + rate_limit_key: " shared-key " + rate_limit_failover_mode: broken + rate_limit_failover_replicas: -1 +`) + + cfg, err := Load(configPath) + if err != nil { + t.Fatalf("Load() error = %v", err) + } + + if cfg.LLM.RequestRateLimitRPM != 0 || + cfg.LLM.RequestRateLimitBurst != 0 || + cfg.LLM.TokenRateLimitTPM != 0 || + cfg.LLM.TokenRateLimitBurst != 0 || + cfg.LLM.EstimatedCharsPerToken != 1 || + cfg.LLM.RateLimitScope != "process" || + cfg.LLM.RateLimitKey != "shared-key" || + cfg.LLM.RateLimitFailoverMode != "fail_closed" || + cfg.LLM.RateLimitFailoverReplicas != 0 { + t.Fatalf("unexpected normalized llm rate limit config: %#v", cfg.LLM) + } +} + +func TestLoadAppliesRabbitMQGenerationPrefetchEnvOverride(t *testing.T) { + t.Setenv("RABBITMQ_GENERATION_CONSUMER_PREFETCH", "1") + + configPath := writeTestConfig(t, ` +rabbitmq: + consumer_prefetch: 10 + generation_consumer_prefetch: 5 +`) + + cfg, err := Load(configPath) + if err != nil { + t.Fatalf("Load() error = %v", err) + } + + if cfg.RabbitMQ.ConsumerPrefetch != 10 || cfg.RabbitMQ.GenerationConsumerPrefetch != 1 { + t.Fatalf("unexpected rabbitmq prefetch config: %#v", cfg.RabbitMQ) + } +} + func TestDatabaseConfigDSNEscapesCredentials(t *testing.T) { cfg := DatabaseConfig{ Host: "db.internal", diff --git a/server/internal/shared/llm/client.go b/server/internal/shared/llm/client.go index 7c4ca03..59e0f4a 100644 --- a/server/internal/shared/llm/client.go +++ b/server/internal/shared/llm/client.go @@ -7,6 +7,8 @@ import ( "strings" "time" + goredis "github.com/redis/go-redis/v9" + "github.com/geo-platform/tenant-api/internal/shared/config" ) @@ -53,15 +55,25 @@ type Client interface { } func New(cfg config.LLMConfig) Client { + return newClient(cfg, nil) +} + +func NewWithRedis(cfg config.LLMConfig, redisClient goredis.Cmdable) Client { + return newClient(cfg, redisClient) +} + +func newClient(cfg config.LLMConfig, redisClient goredis.Cmdable) Client { provider := strings.ToLower(strings.TrimSpace(cfg.Provider)) + var client Client switch provider { case "", "disabled": - return disabledClient{reason: "llm provider is disabled"} + client = disabledClient{reason: "llm provider is disabled"} case "ark": - return NewArkClient(cfg) + client = NewArkClient(cfg) default: - return disabledClient{reason: fmt.Sprintf("unsupported llm provider %q", cfg.Provider)} + client = disabledClient{reason: fmt.Sprintf("unsupported llm provider %q", cfg.Provider)} } + return NewRateLimitedClientWithRedis(client, cfg, redisClient) } type disabledClient struct { diff --git a/server/internal/shared/llm/rate_limited.go b/server/internal/shared/llm/rate_limited.go new file mode 100644 index 0000000..a6116a1 --- /dev/null +++ b/server/internal/shared/llm/rate_limited.go @@ -0,0 +1,472 @@ +package llm + +import ( + "context" + "errors" + "fmt" + "math" + "regexp" + "strings" + "sync/atomic" + "time" + "unicode/utf8" + + goredis "github.com/redis/go-redis/v9" + "golang.org/x/time/rate" + + "github.com/geo-platform/tenant-api/internal/shared/config" +) + +const ( + defaultRateLimitMaxOutputTokens int64 = 4000 + redisRateLimitFailoverCooldown = 5 * time.Second +) + +var redisTokenBucketScript = goredis.NewScript(` +local tokens_key = KEYS[1] +local timestamp_key = KEYS[2] +local cost = tonumber(ARGV[1]) +local rate = tonumber(ARGV[2]) +local capacity = tonumber(ARGV[3]) +local ttl = tonumber(ARGV[4]) + +if cost <= 0 or rate <= 0 or capacity <= 0 then + return 0 +end + +local current_time = redis.call("TIME") +local now = tonumber(current_time[1]) * 1000 + math.floor(tonumber(current_time[2]) / 1000) +local tokens = tonumber(redis.call("GET", tokens_key)) +if tokens == nil then + tokens = capacity +end +local timestamp = tonumber(redis.call("GET", timestamp_key)) +if timestamp == nil then + timestamp = now +end +if now < timestamp then + timestamp = now +end + +local elapsed = now - timestamp +if elapsed > 0 then + tokens = math.min(capacity, tokens + elapsed * rate / 1000) +end + +if tokens >= cost then + tokens = tokens - cost + redis.call("SET", tokens_key, tokens, "PX", ttl) + redis.call("SET", timestamp_key, now, "PX", ttl) + return 0 +end + +local wait = math.ceil((cost - tokens) / rate * 1000) +redis.call("SET", tokens_key, tokens, "PX", ttl) +redis.call("SET", timestamp_key, now, "PX", ttl) +return wait +`) + +type rateLimitedClient struct { + next Client + requestLimiter *rate.Limiter + tokenLimiter *rate.Limiter + redisLimiter *redisRateLimiter + redisFailoverRequests *rate.Limiter + redisFailoverTokens *rate.Limiter + redisUnavailableUntil int64 + rateLimitKey string + requestRateLimitRPM int + requestRateLimitBurst int + tokenRateLimitTPM int64 + tokenRateLimitBurst int + maxOutputTokens int64 + estimatedCharsPerToken float64 +} + +func NewRateLimitedClient(next Client, cfg config.LLMConfig) Client { + return NewRateLimitedClientWithRedis(next, cfg, nil) +} + +func NewRateLimitedClientWithRedis(next Client, cfg config.LLMConfig, redisClient goredis.Cmdable) Client { + if next == nil { + next = disabledClient{reason: "llm client is nil"} + } + + config.NormalizeLLMConfig(&cfg) + if cfg.RequestRateLimitRPM <= 0 && cfg.TokenRateLimitTPM <= 0 { + return next + } + + wrapped := &rateLimitedClient{ + next: next, + rateLimitKey: llmRateLimitKey(cfg), + requestRateLimitRPM: cfg.RequestRateLimitRPM, + requestRateLimitBurst: cfg.RequestRateLimitBurst, + tokenRateLimitTPM: cfg.TokenRateLimitTPM, + tokenRateLimitBurst: cfg.TokenRateLimitBurst, + maxOutputTokens: cfg.MaxOutputTokens, + estimatedCharsPerToken: cfg.EstimatedCharsPerToken, + } + if wrapped.maxOutputTokens <= 0 { + wrapped.maxOutputTokens = defaultRateLimitMaxOutputTokens + } + + if cfg.RateLimitScope == "global" && redisClient != nil { + wrapped.redisLimiter = &redisRateLimiter{redis: redisClient} + wrapped.configureRedisFailover(cfg) + return wrapped + } + + if wrapped.requestRateLimitRPM > 0 { + wrapped.requestLimiter = rate.NewLimiter( + rate.Limit(float64(wrapped.requestRateLimitRPM)/60.0), + rateLimitBurst(wrapped.requestRateLimitRPM, wrapped.requestRateLimitBurst), + ) + } + if wrapped.tokenRateLimitTPM > 0 { + wrapped.tokenLimiter = rate.NewLimiter( + rate.Limit(float64(wrapped.tokenRateLimitTPM)/60.0), + rateLimitInt64Burst(wrapped.tokenRateLimitTPM, wrapped.tokenRateLimitBurst), + ) + } + + return wrapped +} + +func (c *rateLimitedClient) configureRedisFailover(cfg config.LLMConfig) { + if c == nil || cfg.RateLimitFailoverMode != "process_degraded" { + return + } + replicas := cfg.RateLimitFailoverReplicas + if replicas <= 0 { + replicas = 1 + } + if c.requestRateLimitRPM > 0 { + limit := divideLimit(c.requestRateLimitRPM, replicas) + burst := divideLimit(rateLimitBurst(c.requestRateLimitRPM, c.requestRateLimitBurst), replicas) + c.redisFailoverRequests = rate.NewLimiter(rate.Limit(float64(limit)/60.0), burst) + } + if c.tokenRateLimitTPM > 0 { + limit := divideLimit64(c.tokenRateLimitTPM, replicas) + burst := divideLimit(rateLimitInt64Burst(c.tokenRateLimitTPM, c.tokenRateLimitBurst), replicas) + c.redisFailoverTokens = rate.NewLimiter(rate.Limit(float64(limit)/60.0), burst) + } +} + +func (c *rateLimitedClient) Validate() error { + if c == nil || c.next == nil { + return disabledClient{reason: "llm client is nil"}.Validate() + } + return c.next.Validate() +} + +func (c *rateLimitedClient) Generate(ctx context.Context, req GenerateRequest, onDelta func(string)) (*GenerateResult, error) { + if err := c.Validate(); err != nil { + return nil, err + } + tokens := estimateRateLimitTokens(req, c.maxOutputTokens, c.estimatedCharsPerToken) + if c.redisLimiter != nil { + if err := c.waitGlobalRateLimit(ctx, "requests", 1, int64(c.requestRateLimitRPM), c.requestRateLimitBurst); err != nil { + return nil, err + } + if err := c.waitGlobalRateLimit(ctx, "tokens", tokens, c.tokenRateLimitTPM, c.tokenRateLimitBurst); err != nil { + return nil, err + } + return c.next.Generate(ctx, req, onDelta) + } + if c.requestLimiter != nil { + if err := c.requestLimiter.Wait(ctx); err != nil { + return nil, fmt.Errorf("llm request rate limit wait: %w", err) + } + } + if c.tokenLimiter != nil { + if err := waitLimiterN(ctx, c.tokenLimiter, tokens); err != nil { + return nil, fmt.Errorf("llm token rate limit wait: %w", err) + } + } + return c.next.Generate(ctx, req, onDelta) +} + +func (c *rateLimitedClient) waitGlobalRateLimit(ctx context.Context, bucket string, cost int, limitPerMinute int64, burst int) error { + if c == nil || c.redisLimiter == nil || cost <= 0 || limitPerMinute <= 0 { + return nil + } + if c.redisCircuitOpen(time.Now()) { + return c.waitRedisFailoverRateLimit(ctx, bucket, cost) + } + if err := c.waitRedisRateLimit(ctx, bucket, cost, limitPerMinute, burst); err != nil { + if !c.canUseRedisFailover(err) { + return err + } + c.markRedisUnavailable(time.Now()) + if fallbackErr := c.waitRedisFailoverRateLimit(ctx, bucket, cost); fallbackErr != nil { + return fmt.Errorf("llm %s global rate limit failed and process-degraded failover also failed: %w", bucket, errors.Join(err, fallbackErr)) + } + return nil + } + return nil +} + +func (c *rateLimitedClient) canUseRedisFailover(err error) bool { + if c == nil || err == nil { + return false + } + if c.redisFailoverRequests == nil && c.redisFailoverTokens == nil { + return false + } + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return false + } + return true +} + +func (c *rateLimitedClient) waitRedisFailoverRateLimit(ctx context.Context, bucket string, cost int) error { + switch bucket { + case "requests": + if c.redisFailoverRequests == nil { + return nil + } + if err := c.redisFailoverRequests.Wait(ctx); err != nil { + return fmt.Errorf("llm %s process-degraded rate limit wait: %w", bucket, err) + } + case "tokens": + if c.redisFailoverTokens == nil { + return nil + } + if err := waitLimiterN(ctx, c.redisFailoverTokens, cost); err != nil { + return fmt.Errorf("llm %s process-degraded rate limit wait: %w", bucket, err) + } + } + return nil +} + +func (c *rateLimitedClient) redisCircuitOpen(now time.Time) bool { + if c == nil { + return false + } + until := atomic.LoadInt64(&c.redisUnavailableUntil) + return until > now.UnixNano() +} + +func (c *rateLimitedClient) markRedisUnavailable(now time.Time) { + if c == nil { + return + } + atomic.StoreInt64(&c.redisUnavailableUntil, now.Add(redisRateLimitFailoverCooldown).UnixNano()) +} + +func (c *rateLimitedClient) waitRedisRateLimit(ctx context.Context, bucket string, cost int, limitPerMinute int64, burst int) error { + if c == nil || c.redisLimiter == nil || cost <= 0 || limitPerMinute <= 0 { + return nil + } + capacity := rateLimitInt64Burst(limitPerMinute, burst) + if capacity <= 0 { + return nil + } + remaining := cost + for remaining > 0 { + chunk := remaining + if chunk > capacity { + chunk = capacity + } + if err := c.redisLimiter.Wait(ctx, c.rateLimitKey, bucket, chunk, limitPerMinute, capacity); err != nil { + return fmt.Errorf("llm %s global rate limit wait: %w", bucket, err) + } + remaining -= chunk + } + return nil +} + +func estimateRateLimitTokens(req GenerateRequest, fallbackMaxOutputTokens int64, charsPerToken float64) int { + if charsPerToken <= 0 { + charsPerToken = 1 + } + promptChars := utf8.RuneCountInString(strings.TrimSpace(req.Prompt)) + promptTokens := int64(math.Ceil(float64(promptChars) / charsPerToken)) + outputTokens := req.MaxOutputTokens + if outputTokens <= 0 { + outputTokens = fallbackMaxOutputTokens + } + if outputTokens <= 0 { + outputTokens = defaultRateLimitMaxOutputTokens + } + + total := promptTokens + outputTokens + if total <= 0 { + return 1 + } + if total > int64(maxInt()) { + return maxInt() + } + return int(total) +} + +func waitLimiterN(ctx context.Context, limiter *rate.Limiter, n int) error { + if limiter == nil || n <= 0 { + return nil + } + burst := limiter.Burst() + if burst <= 0 { + return nil + } + for n > 0 { + chunk := n + if chunk > burst { + chunk = burst + } + if err := limiter.WaitN(ctx, chunk); err != nil { + return err + } + n -= chunk + } + return nil +} + +func rateLimitBurst(limit, configured int) int { + if configured > 0 { + return configured + } + burst := limit / 10 + if burst < 1 { + return 1 + } + return burst +} + +func rateLimitInt64Burst(limit int64, configured int) int { + if configured > 0 { + return configured + } + burst := limit / 10 + if burst < 1 { + return 1 + } + if burst > int64(maxInt()) { + return maxInt() + } + return int(burst) +} + +func divideLimit(limit, divisor int) int { + if limit <= 0 { + return 0 + } + if divisor <= 1 { + return limit + } + quotient := limit / divisor + if quotient < 1 { + return 1 + } + return quotient +} + +func divideLimit64(limit int64, divisor int) int64 { + if limit <= 0 { + return 0 + } + if divisor <= 1 { + return limit + } + quotient := limit / int64(divisor) + if quotient < 1 { + return 1 + } + return quotient +} + +func maxInt() int { + return int(^uint(0) >> 1) +} + +type redisRateLimiter struct { + redis goredis.Cmdable +} + +func (l *redisRateLimiter) Wait(ctx context.Context, group, bucket string, cost int, limitPerMinute int64, capacity int) error { + if l == nil || l.redis == nil || cost <= 0 || limitPerMinute <= 0 || capacity <= 0 { + return nil + } + + ratePerSecond := float64(limitPerMinute) / 60.0 + if ratePerSecond <= 0 { + return nil + } + ttl := redisBucketTTL(ratePerSecond, capacity) + keys := redisBucketKeys(group, bucket) + + for { + waitMillis, err := redisTokenBucketScript.Run( + ctx, + l.redis, + keys, + cost, + ratePerSecond, + capacity, + ttl.Milliseconds(), + ).Int64() + if err != nil { + return err + } + if waitMillis <= 0 { + return nil + } + + timer := time.NewTimer(time.Duration(waitMillis) * time.Millisecond) + select { + case <-ctx.Done(): + timer.Stop() + return ctx.Err() + case <-timer.C: + } + } +} + +func redisBucketTTL(ratePerSecond float64, capacity int) time.Duration { + if ratePerSecond <= 0 || capacity <= 0 { + return 2 * time.Minute + } + refill := time.Duration(math.Ceil(float64(capacity)/ratePerSecond)) * time.Second + if refill < time.Minute { + return 2 * time.Minute + } + return refill * 2 +} + +func redisBucketKeys(group, bucket string) []string { + escapedGroup := sanitizeRateLimitKey(group) + escapedBucket := sanitizeRateLimitKey(bucket) + prefix := fmt.Sprintf("geo:llm:rate:{%s}:%s", escapedGroup, escapedBucket) + return []string{prefix + ":tokens", prefix + ":ts"} +} + +func llmRateLimitKey(cfg config.LLMConfig) string { + if key := strings.TrimSpace(cfg.RateLimitKey); key != "" { + return key + } + parts := []string{ + strings.TrimSpace(cfg.Provider), + strings.TrimSpace(cfg.BaseURL), + strings.TrimSpace(cfg.Model), + } + key := strings.Join(parts, ":") + if strings.Trim(key, ": ") == "" { + return "default" + } + return key +} + +var rateLimitKeySanitizer = regexp.MustCompile(`[^a-zA-Z0-9_.:-]+`) + +func sanitizeRateLimitKey(value string) string { + value = strings.TrimSpace(value) + if value == "" { + return "default" + } + value = rateLimitKeySanitizer.ReplaceAllString(value, "_") + value = strings.Trim(value, "_") + if value == "" { + return "default" + } + return value +} diff --git a/server/internal/shared/llm/rate_limited_test.go b/server/internal/shared/llm/rate_limited_test.go new file mode 100644 index 0000000..47fb622 --- /dev/null +++ b/server/internal/shared/llm/rate_limited_test.go @@ -0,0 +1,256 @@ +package llm + +import ( + "context" + "testing" + "time" + + "github.com/alicebob/miniredis/v2" + goredis "github.com/redis/go-redis/v9" + + "github.com/geo-platform/tenant-api/internal/shared/config" +) + +type stubLLMClient struct { + req GenerateRequest + calls int +} + +func (c *stubLLMClient) Validate() error { + return nil +} + +func (c *stubLLMClient) Generate(_ context.Context, req GenerateRequest, _ func(string)) (*GenerateResult, error) { + c.req = req + c.calls++ + return &GenerateResult{Content: "ok", Model: req.Model}, nil +} + +func TestNewRateLimitedClientReturnsBaseWhenLimitsDisabled(t *testing.T) { + base := &stubLLMClient{} + got := NewRateLimitedClient(base, config.LLMConfig{}) + if got != base { + t.Fatalf("expected base client when rate limits are disabled") + } +} + +func TestRateLimitedClientDelegatesGenerate(t *testing.T) { + base := &stubLLMClient{} + client := NewRateLimitedClient(base, config.LLMConfig{ + RequestRateLimitRPM: 60000, + TokenRateLimitTPM: 600000, + MaxOutputTokens: 100, + EstimatedCharsPerToken: 1, + }) + + result, err := client.Generate(context.Background(), GenerateRequest{ + Prompt: "hello", + MaxOutputTokens: 10, + }, nil) + if err != nil { + t.Fatalf("Generate() error = %v", err) + } + if result.Content != "ok" { + t.Fatalf("Generate() content = %q, want ok", result.Content) + } + if base.req.Prompt != "hello" { + t.Fatalf("delegated prompt = %q, want hello", base.req.Prompt) + } + if base.calls != 1 { + t.Fatalf("delegated calls = %d, want 1", base.calls) + } +} + +func TestEstimateRateLimitTokensIncludesFallbackOutput(t *testing.T) { + got := estimateRateLimitTokens(GenerateRequest{Prompt: " 你好世界 "}, 16000, 1) + if got != 16004 { + t.Fatalf("estimateRateLimitTokens() = %d, want 16004", got) + } +} + +func TestEstimateRateLimitTokensUsesCharsPerToken(t *testing.T) { + got := estimateRateLimitTokens(GenerateRequest{ + Prompt: "abcd", + MaxOutputTokens: 6, + }, 16000, 2) + if got != 8 { + t.Fatalf("estimateRateLimitTokens() = %d, want 8", got) + } +} + +func TestRateLimitedClientStopsBeforeDelegateWhenContextCanceled(t *testing.T) { + base := &stubLLMClient{} + client := NewRateLimitedClient(base, config.LLMConfig{ + RequestRateLimitRPM: 1, + RequestRateLimitBurst: 1, + TokenRateLimitTPM: 1, + TokenRateLimitBurst: 1, + MaxOutputTokens: 100, + EstimatedCharsPerToken: 1, + }) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + if _, err := client.Generate(ctx, GenerateRequest{Prompt: "hello"}, nil); err == nil { + t.Fatalf("Generate() expected context cancellation error") + } + if base.calls != 0 { + t.Fatalf("delegated calls = %d, want 0", base.calls) + } +} + +func TestGlobalRateLimitedClientsShareRedisBudget(t *testing.T) { + mr := miniredis.RunT(t) + rdb := goredis.NewClient(&goredis.Options{Addr: mr.Addr()}) + defer rdb.Close() + + cfg := config.LLMConfig{ + RateLimitScope: "global", + RateLimitKey: "ark-shared", + RequestRateLimitRPM: 60, + RequestRateLimitBurst: 1, + } + firstBase := &stubLLMClient{} + first := NewRateLimitedClientWithRedis(firstBase, cfg, rdb) + if _, err := first.Generate(context.Background(), GenerateRequest{Prompt: "first"}, nil); err != nil { + t.Fatalf("first Generate() error = %v", err) + } + if firstBase.calls != 1 { + t.Fatalf("first delegated calls = %d, want 1", firstBase.calls) + } + + secondBase := &stubLLMClient{} + second := NewRateLimitedClientWithRedis(secondBase, cfg, rdb) + ctx, cancel := context.WithTimeout(context.Background(), 25*time.Millisecond) + defer cancel() + + if _, err := second.Generate(ctx, GenerateRequest{Prompt: "second"}, nil); err == nil { + t.Fatalf("second Generate() expected global rate limit wait error") + } + if secondBase.calls != 0 { + t.Fatalf("second delegated calls = %d, want 0", secondBase.calls) + } +} + +func TestGlobalRateLimitedClientsShareRedisTokenBudget(t *testing.T) { + mr := miniredis.RunT(t) + rdb := goredis.NewClient(&goredis.Options{Addr: mr.Addr()}) + defer rdb.Close() + + cfg := config.LLMConfig{ + RateLimitScope: "global", + RateLimitKey: "ark-token-shared", + TokenRateLimitTPM: 60, + TokenRateLimitBurst: 10, + MaxOutputTokens: 5, + EstimatedCharsPerToken: 1, + } + firstBase := &stubLLMClient{} + first := NewRateLimitedClientWithRedis(firstBase, cfg, rdb) + if _, err := first.Generate(context.Background(), GenerateRequest{ + Prompt: "12345", + MaxOutputTokens: 5, + }, nil); err != nil { + t.Fatalf("first Generate() error = %v", err) + } + if firstBase.calls != 1 { + t.Fatalf("first delegated calls = %d, want 1", firstBase.calls) + } + + secondBase := &stubLLMClient{} + second := NewRateLimitedClientWithRedis(secondBase, cfg, rdb) + ctx, cancel := context.WithTimeout(context.Background(), 25*time.Millisecond) + defer cancel() + + if _, err := second.Generate(ctx, GenerateRequest{ + Prompt: "12345", + MaxOutputTokens: 5, + }, nil); err == nil { + t.Fatalf("second Generate() expected global token rate limit wait error") + } + if secondBase.calls != 0 { + t.Fatalf("second delegated calls = %d, want 0", secondBase.calls) + } +} + +func TestGlobalRateLimitedClientFailsClosedWhenRedisUnavailable(t *testing.T) { + mr := miniredis.RunT(t) + addr := mr.Addr() + mr.Close() + + rdb := goredis.NewClient(&goredis.Options{ + Addr: addr, + DialTimeout: 10 * time.Millisecond, + ReadTimeout: 10 * time.Millisecond, + WriteTimeout: 10 * time.Millisecond, + }) + defer rdb.Close() + + cfg := config.LLMConfig{ + RateLimitScope: "global", + RateLimitKey: "ark-fail-closed", + RequestRateLimitRPM: 60, + RequestRateLimitBurst: 1, + } + base := &stubLLMClient{} + client := NewRateLimitedClientWithRedis(base, cfg, rdb) + + if _, err := client.Generate(context.Background(), GenerateRequest{Prompt: "blocked"}, nil); err == nil { + t.Fatalf("Generate() expected redis rate limit error") + } + if base.calls != 0 { + t.Fatalf("delegated calls = %d, want 0", base.calls) + } +} + +func TestGlobalRateLimitedClientUsesProcessDegradedFailoverWhenRedisUnavailable(t *testing.T) { + mr := miniredis.RunT(t) + addr := mr.Addr() + mr.Close() + + rdb := goredis.NewClient(&goredis.Options{ + Addr: addr, + DialTimeout: 10 * time.Millisecond, + ReadTimeout: 10 * time.Millisecond, + WriteTimeout: 10 * time.Millisecond, + }) + defer rdb.Close() + + cfg := config.LLMConfig{ + RateLimitScope: "global", + RateLimitKey: "ark-process-degraded", + RateLimitFailoverMode: "process_degraded", + RateLimitFailoverReplicas: 3, + RequestRateLimitRPM: 60, + RequestRateLimitBurst: 3, + TokenRateLimitTPM: 60, + TokenRateLimitBurst: 3, + MaxOutputTokens: 1, + EstimatedCharsPerToken: 1, + } + base := &stubLLMClient{} + client := NewRateLimitedClientWithRedis(base, cfg, rdb) + + if _, err := client.Generate(context.Background(), GenerateRequest{ + Prompt: "", + MaxOutputTokens: 1, + }, nil); err != nil { + t.Fatalf("first Generate() error = %v", err) + } + if base.calls != 1 { + t.Fatalf("delegated calls = %d, want 1", base.calls) + } + + ctx, cancel := context.WithTimeout(context.Background(), 25*time.Millisecond) + defer cancel() + if _, err := client.Generate(ctx, GenerateRequest{ + Prompt: "", + MaxOutputTokens: 1, + }, nil); err == nil { + t.Fatalf("second Generate() expected process-degraded rate limit wait error") + } + if base.calls != 1 { + t.Fatalf("delegated calls after blocked call = %d, want 1", base.calls) + } +} diff --git a/server/internal/shared/messaging/rabbitmq/client.go b/server/internal/shared/messaging/rabbitmq/client.go index 756115a..2f08f3f 100644 --- a/server/internal/shared/messaging/rabbitmq/client.go +++ b/server/internal/shared/messaging/rabbitmq/client.go @@ -280,7 +280,11 @@ func (c *Client) ConsumeMonitorProjectionRebuild(consumerName string) (<-chan am } func (c *Client) ConsumeGenerationTask(consumerName string) (<-chan amqp.Delivery, *amqp.Channel, error) { - return c.consume(consumerName, c.cfg.GenerationQueue) + prefetch := c.cfg.GenerationConsumerPrefetch + if prefetch <= 0 { + prefetch = c.cfg.ConsumerPrefetch + } + return c.consumeWithPrefetch(consumerName, c.cfg.GenerationQueue, prefetch) } func (c *Client) ConsumeGenerationEvents(consumerName string) (<-chan amqp.Delivery, *amqp.Channel, error) { @@ -312,14 +316,18 @@ func (c *Client) ConsumeComplianceReviewTask(consumerName string) (<-chan amqp.D } func (c *Client) consume(consumerName, queue string) (<-chan amqp.Delivery, *amqp.Channel, error) { + return c.consumeWithPrefetch(consumerName, queue, c.cfg.ConsumerPrefetch) +} + +func (c *Client) consumeWithPrefetch(consumerName, queue string, prefetch int) (<-chan amqp.Delivery, *amqp.Channel, error) { for attempt := 0; attempt < 2; attempt++ { conn, ch, err := c.openChannel() if err != nil { return nil, nil, err } - if c.cfg.ConsumerPrefetch > 0 { - if err := ch.Qos(c.cfg.ConsumerPrefetch, 0, false); err != nil { + if prefetch > 0 { + if err := ch.Qos(prefetch, 0, false); err != nil { _ = ch.Close() c.invalidateConnection(conn) if attempt == 1 { diff --git a/server/internal/shared/middleware/logger_test.go b/server/internal/shared/middleware/logger_test.go index 0101c0a..bd6e50b 100644 --- a/server/internal/shared/middleware/logger_test.go +++ b/server/internal/shared/middleware/logger_test.go @@ -18,8 +18,6 @@ import ( ) func TestLoggerIncludesDetailedAuthFailureContext(t *testing.T) { - t.Parallel() - gin.SetMode(gin.TestMode) var logBuffer bytes.Buffer @@ -74,8 +72,6 @@ func TestLoggerIncludesDetailedAuthFailureContext(t *testing.T) { } func TestLoggerUsesForwardedClientIPFromTrustedProxy(t *testing.T) { - t.Parallel() - gin.SetMode(gin.TestMode) var logBuffer bytes.Buffer @@ -107,8 +103,6 @@ func TestLoggerUsesForwardedClientIPFromTrustedProxy(t *testing.T) { } func TestLoggerRedactsSensitiveQueryValues(t *testing.T) { - t.Parallel() - gin.SetMode(gin.TestMode) var logBuffer bytes.Buffer diff --git a/server/internal/shared/middleware/security_headers_test.go b/server/internal/shared/middleware/security_headers_test.go index 130c95e..33653b5 100644 --- a/server/internal/shared/middleware/security_headers_test.go +++ b/server/internal/shared/middleware/security_headers_test.go @@ -10,8 +10,6 @@ import ( ) func TestSecurityHeadersAddsNoStoreForSensitiveAuthPaths(t *testing.T) { - t.Parallel() - gin.SetMode(gin.TestMode) router := gin.New() router.Use(SecurityHeaders()) @@ -34,8 +32,6 @@ func TestSecurityHeadersAddsNoStoreForSensitiveAuthPaths(t *testing.T) { } func TestCORSWithConfigRejectsUnknownPreflightOrigin(t *testing.T) { - t.Parallel() - gin.SetMode(gin.TestMode) router := gin.New() router.Use(CORSWithConfig(CORSConfig{AllowedOrigins: []string{"https://admin.example.com"}})) @@ -54,8 +50,6 @@ func TestCORSWithConfigRejectsUnknownPreflightOrigin(t *testing.T) { } func TestCORSWithConfigAllowsConfiguredOrigin(t *testing.T) { - t.Parallel() - gin.SetMode(gin.TestMode) router := gin.New() router.Use(CORSWithConfig(CORSConfig{AllowedOrigins: []string{"https://admin.example.com"}}))