geo/server/internal/shared/llm/rate_limited_test.go

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)
	}
}