feat(llm): add configurable request/token rate limiting for LLM client

Introduce process-local and Redis-backed global rate limiting for the shared LLM client, plus a dedicated generation consumer prefetch knob. - Add request (RPM) and token (TPM) rate limits with burst controls, process or global scope, and fail-closed/process-degraded failover - Wire the LLM client through NewWithRedis so limits can be shared across replicas via a Redis token bucket - Add generation_consumer_prefetch to tune RabbitMQ generation consumers independently of the default prefetch - Bump generation worker_concurrency to 24 and promote golang.org/x/time to a direct dependency - Drop t.Parallel() from a few middleware/bootstrap tests Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-10 00:14:40 +08:00
parent 0acd37918e
commit 347cea1296
16 changed files with 1016 additions and 30 deletions
@@ -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",